mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 23:58:11 +00:00
removed anime4k support and added Anime4KCPP support; cleaned up the program
This commit is contained in:
parent
0f1639ed62
commit
88299d404a
@ -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: April 26, 2020
|
Last Modified: May 3, 2020
|
||||||
|
|
||||||
Description: This file contains the Upscaler class. Each
|
Description: This file contains the Upscaler class. Each
|
||||||
instance of the Upscaler class is an upscaler on an image or
|
instance of the Upscaler class is an upscaler on an image or
|
||||||
@ -12,19 +12,20 @@ a folder.
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
# local imports
|
# local imports
|
||||||
from anime4k import Anime4k
|
|
||||||
from exceptions import *
|
from exceptions import *
|
||||||
from ffmpeg import Ffmpeg
|
|
||||||
from image_cleaner import ImageCleaner
|
from image_cleaner import ImageCleaner
|
||||||
from srmd_ncnn_vulkan import SrmdNcnnVulkan
|
from wrappers.anime4kcpp import Anime4kCpp
|
||||||
from waifu2x_caffe import Waifu2xCaffe
|
from wrappers.ffmpeg import Ffmpeg
|
||||||
from waifu2x_converter import Waifu2xConverter
|
from wrappers.srmd_ncnn_vulkan import SrmdNcnnVulkan
|
||||||
from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
|
from wrappers.waifu2x_caffe import Waifu2xCaffe
|
||||||
|
from wrappers.waifu2x_converter_cpp import Waifu2xConverterCpp
|
||||||
|
from wrappers.waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
|
||||||
|
|
||||||
# built-in imports
|
# built-in imports
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
import contextlib
|
import contextlib
|
||||||
import copy
|
import copy
|
||||||
|
import os
|
||||||
import pathlib
|
import pathlib
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
@ -38,7 +39,7 @@ import traceback
|
|||||||
from avalon_framework import Avalon
|
from avalon_framework import Avalon
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k', 'srmd_ncnn_vulkan']
|
AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan', 'anime4kcpp']
|
||||||
|
|
||||||
|
|
||||||
class Upscaler:
|
class Upscaler:
|
||||||
@ -50,16 +51,15 @@ class Upscaler:
|
|||||||
ArgumentError -- if argument is not valid
|
ArgumentError -- if argument is not valid
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, input_video, output_video, method, driver_settings, ffmpeg_settings):
|
def __init__(self, input_video, output_video, driver_settings, ffmpeg_settings):
|
||||||
# mandatory arguments
|
# mandatory arguments
|
||||||
self.input_video = input_video
|
self.input_video = input_video
|
||||||
self.output_video = output_video
|
self.output_video = output_video
|
||||||
self.method = method
|
|
||||||
self.driver_settings = driver_settings
|
self.driver_settings = driver_settings
|
||||||
self.ffmpeg_settings = ffmpeg_settings
|
self.ffmpeg_settings = ffmpeg_settings
|
||||||
|
|
||||||
# optional arguments
|
# optional arguments
|
||||||
self.waifu2x_driver = 'waifu2x_caffe'
|
self.driver = 'waifu2x_caffe'
|
||||||
self.scale_width = None
|
self.scale_width = None
|
||||||
self.scale_height = None
|
self.scale_height = None
|
||||||
self.scale_ratio = None
|
self.scale_ratio = None
|
||||||
@ -101,13 +101,26 @@ class Upscaler:
|
|||||||
# check if arguments are valid / all necessary argument
|
# check if arguments are valid / all necessary argument
|
||||||
# values are specified
|
# values are specified
|
||||||
if not self.input_video:
|
if not self.input_video:
|
||||||
raise ArgumentError('You need to specify the video to process')
|
Avalon.error('You must specify input video file/directory path')
|
||||||
elif (not self.scale_width or not self.scale_height) and not self.scale_ratio:
|
raise ArgumentError('input video path not specified')
|
||||||
raise ArgumentError('You must specify output video width and height or upscale factor')
|
if not self.output_video:
|
||||||
elif not self.output_video:
|
Avalon.error('You must specify output video file/directory path')
|
||||||
raise ArgumentError('You need to specify the output video name')
|
raise ArgumentError('output video path not specified')
|
||||||
elif not self.method:
|
if (self.driver in ['waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k']) and self.scale_width and self.scale_height:
|
||||||
raise ArgumentError('You need to specify the enlarging processing unit')
|
Avalon.error('Selected driver accepts only scaling ratio')
|
||||||
|
raise ArgumentError('selected driver supports only scaling ratio')
|
||||||
|
if self.driver == 'waifu2x_ncnn_vulkan' and self.scale_ratio is not None and (self.scale_ratio > 2 or not self.scale_ratio.is_integer()):
|
||||||
|
Avalon.error('Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
|
||||||
|
raise ArgumentError('scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
|
||||||
|
if self.driver == 'srmd_ncnn_vulkan' and self.scale_ratio is not None and (self.scale_ratio not in [2, 3, 4]):
|
||||||
|
Avalon.error('Scaling ratio must be one of 2, 3 or 4 for srmd_ncnn_vulkan')
|
||||||
|
raise ArgumentError('scaling ratio must be one of 2, 3 or 4 for srmd_ncnn_vulkan')
|
||||||
|
if (self.scale_width or self.scale_height) and self.scale_ratio:
|
||||||
|
Avalon.error('You can only specify either scaling ratio or output width and height')
|
||||||
|
raise ArgumentError('both scaling ration and width/height specified')
|
||||||
|
if (self.scale_width and not self.scale_height) or (not self.scale_width and self.scale_height):
|
||||||
|
Avalon.error('You must specify both width and height')
|
||||||
|
raise ArgumentError('only one of width or height is specified')
|
||||||
|
|
||||||
def _progress_bar(self, extracted_frames_directories):
|
def _progress_bar(self, extracted_frames_directories):
|
||||||
""" This method prints a progress bar
|
""" This method prints a progress bar
|
||||||
@ -159,8 +172,8 @@ class Upscaler:
|
|||||||
self.progress_bar_exit_signal = False
|
self.progress_bar_exit_signal = False
|
||||||
|
|
||||||
# initialize waifu2x driver
|
# initialize waifu2x driver
|
||||||
if self.waifu2x_driver not in AVAILABLE_DRIVERS:
|
if self.driver not in AVAILABLE_DRIVERS:
|
||||||
raise UnrecognizedDriverError(f'Unrecognized driver: {self.waifu2x_driver}')
|
raise UnrecognizedDriverError(f'Unrecognized driver: {self.driver}')
|
||||||
|
|
||||||
# create a container for all upscaler processes
|
# create a container for all upscaler processes
|
||||||
upscaler_processes = []
|
upscaler_processes = []
|
||||||
@ -186,7 +199,7 @@ class Upscaler:
|
|||||||
process_directory.mkdir(parents=True, exist_ok=True)
|
process_directory.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
# waifu2x-converter-cpp will perform multi-threading within its own process
|
# waifu2x-converter-cpp will perform multi-threading within its own process
|
||||||
if self.waifu2x_driver in ['waifu2x_converter', 'anime4k'] :
|
if self.driver == 'waifu2x_converter_cpp':
|
||||||
process_directories = [self.extracted_frames]
|
process_directories = [self.extracted_frames]
|
||||||
|
|
||||||
else:
|
else:
|
||||||
@ -202,8 +215,8 @@ class Upscaler:
|
|||||||
for process_directory in process_directories:
|
for process_directory in process_directories:
|
||||||
|
|
||||||
# if the driver being used is waifu2x-caffe
|
# if the driver being used is waifu2x-caffe
|
||||||
if self.waifu2x_driver == 'waifu2x_caffe':
|
if self.driver == 'waifu2x_caffe':
|
||||||
driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.method, self.model_dir, self.bit_depth)
|
driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.model_dir, self.bit_depth)
|
||||||
if self.scale_ratio:
|
if self.scale_ratio:
|
||||||
upscaler_processes.append(driver.upscale(process_directory,
|
upscaler_processes.append(driver.upscale(process_directory,
|
||||||
self.upscaled_frames,
|
self.upscaled_frames,
|
||||||
@ -220,8 +233,8 @@ class Upscaler:
|
|||||||
self.image_format))
|
self.image_format))
|
||||||
|
|
||||||
# if the driver being used is waifu2x-converter-cpp
|
# if the driver being used is waifu2x-converter-cpp
|
||||||
elif self.waifu2x_driver == 'waifu2x_converter':
|
elif self.driver == 'waifu2x_converter_cpp':
|
||||||
driver = Waifu2xConverter(self.driver_settings, self.model_dir)
|
driver = Waifu2xConverterCpp(self.driver_settings, self.model_dir)
|
||||||
upscaler_processes.append(driver.upscale(process_directory,
|
upscaler_processes.append(driver.upscale(process_directory,
|
||||||
self.upscaled_frames,
|
self.upscaled_frames,
|
||||||
self.scale_ratio,
|
self.scale_ratio,
|
||||||
@ -229,22 +242,14 @@ class Upscaler:
|
|||||||
self.image_format))
|
self.image_format))
|
||||||
|
|
||||||
# if the driver being used is waifu2x-ncnn-vulkan
|
# if the driver being used is waifu2x-ncnn-vulkan
|
||||||
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
|
elif self.driver == 'waifu2x_ncnn_vulkan':
|
||||||
driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings))
|
driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings))
|
||||||
upscaler_processes.append(driver.upscale(process_directory,
|
upscaler_processes.append(driver.upscale(process_directory,
|
||||||
self.upscaled_frames,
|
self.upscaled_frames,
|
||||||
self.scale_ratio))
|
self.scale_ratio))
|
||||||
|
|
||||||
# if the driver being used is anime4k
|
|
||||||
elif self.waifu2x_driver == 'anime4k':
|
|
||||||
driver = Anime4k(copy.deepcopy(self.driver_settings))
|
|
||||||
upscaler_processes += driver.upscale(process_directory,
|
|
||||||
self.upscaled_frames,
|
|
||||||
self.scale_ratio,
|
|
||||||
self.processes)
|
|
||||||
|
|
||||||
# if the driver being used is srmd_ncnn_vulkan
|
# if the driver being used is srmd_ncnn_vulkan
|
||||||
elif self.waifu2x_driver == 'srmd_ncnn_vulkan':
|
elif self.driver == 'srmd_ncnn_vulkan':
|
||||||
driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings))
|
driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings))
|
||||||
upscaler_processes.append(driver.upscale(process_directory,
|
upscaler_processes.append(driver.upscale(process_directory,
|
||||||
self.upscaled_frames,
|
self.upscaled_frames,
|
||||||
@ -278,7 +283,7 @@ class Upscaler:
|
|||||||
|
|
||||||
# if the driver is waifu2x-converter-cpp
|
# if the driver is waifu2x-converter-cpp
|
||||||
# images need to be renamed to be recognizable for FFmpeg
|
# images need to be renamed to be recognizable for FFmpeg
|
||||||
if self.waifu2x_driver == 'waifu2x_converter':
|
if self.driver == 'waifu2x_converter_cpp':
|
||||||
for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]:
|
for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]:
|
||||||
renamed = re.sub(f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.image_format}', f'.{self.image_format}', str(image.name))
|
renamed = re.sub(f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.image_format}', f'.{self.image_format}', str(image.name))
|
||||||
(self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
|
(self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
|
||||||
@ -305,12 +310,22 @@ class Upscaler:
|
|||||||
self.input_video = self.input_video.absolute()
|
self.input_video = self.input_video.absolute()
|
||||||
self.output_video = self.output_video.absolute()
|
self.output_video = self.output_video.absolute()
|
||||||
|
|
||||||
|
# drivers that have native support for video processing
|
||||||
|
if self.driver == 'anime4kcpp':
|
||||||
|
# append FFmpeg path to the end of PATH
|
||||||
|
# Anime4KCPP will then use FFmpeg to migrate audio tracks
|
||||||
|
os.environ['PATH'] += f';{self.ffmpeg_settings["ffmpeg_path"]}'
|
||||||
|
Avalon.info('Starting to upscale extracted images')
|
||||||
|
driver = Anime4kCpp(self.driver_settings)
|
||||||
|
driver.upscale(self.input_video, self.output_video, self.scale_ratio, self.processes).wait()
|
||||||
|
Avalon.info('Upscaling completed')
|
||||||
|
|
||||||
|
else:
|
||||||
|
self.create_temp_directories()
|
||||||
|
|
||||||
# initialize objects for ffmpeg and waifu2x-caffe
|
# initialize objects for ffmpeg and waifu2x-caffe
|
||||||
fm = Ffmpeg(self.ffmpeg_settings, self.image_format)
|
fm = Ffmpeg(self.ffmpeg_settings, self.image_format)
|
||||||
|
|
||||||
# extract frames from video
|
|
||||||
fm.extract_frames(self.input_video, self.extracted_frames)
|
|
||||||
|
|
||||||
Avalon.info('Reading video information')
|
Avalon.info('Reading video information')
|
||||||
video_info = fm.get_video_info(self.input_video)
|
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
|
||||||
@ -328,6 +343,9 @@ class Upscaler:
|
|||||||
Avalon.error('Aborting: No video stream found')
|
Avalon.error('Aborting: No video stream found')
|
||||||
raise StreamNotFoundError('no video stream found')
|
raise StreamNotFoundError('no video stream found')
|
||||||
|
|
||||||
|
# extract frames from video
|
||||||
|
fm.extract_frames(self.input_video, self.extracted_frames)
|
||||||
|
|
||||||
# get average frame rate of video stream
|
# get average frame rate of video stream
|
||||||
framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
|
framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
|
||||||
fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']
|
fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']
|
||||||
@ -366,3 +384,6 @@ class Upscaler:
|
|||||||
# migrate audio tracks and subtitles
|
# migrate audio tracks and subtitles
|
||||||
Avalon.info('Migrating audio tracks and subtitles to upscaled video')
|
Avalon.info('Migrating audio tracks and subtitles to upscaled video')
|
||||||
fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)
|
fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)
|
||||||
|
|
||||||
|
# destroy temp directories
|
||||||
|
self.cleanup_temp_directories()
|
||||||
|
210
src/video2x.py
210
src/video2x.py
@ -13,7 +13,7 @@ __ __ _ _ ___ __ __
|
|||||||
Name: Video2X Controller
|
Name: Video2X Controller
|
||||||
Creator: K4YT3X
|
Creator: K4YT3X
|
||||||
Date Created: Feb 24, 2018
|
Date Created: Feb 24, 2018
|
||||||
Last Modified: April 26, 2020
|
Last Modified: May 3, 2020
|
||||||
|
|
||||||
Editor: BrianPetkovsek
|
Editor: BrianPetkovsek
|
||||||
Last Modified: June 17, 2019
|
Last Modified: June 17, 2019
|
||||||
@ -49,7 +49,6 @@ smooth and edges sharp.
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
# local imports
|
# local imports
|
||||||
from exceptions import *
|
|
||||||
from upscaler import AVAILABLE_DRIVERS
|
from upscaler import AVAILABLE_DRIVERS
|
||||||
from upscaler import Upscaler
|
from upscaler import Upscaler
|
||||||
|
|
||||||
@ -59,7 +58,6 @@ import contextlib
|
|||||||
import pathlib
|
import pathlib
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
import subprocess
|
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import time
|
import time
|
||||||
@ -68,11 +66,9 @@ import yaml
|
|||||||
|
|
||||||
# third-party imports
|
# third-party imports
|
||||||
from avalon_framework import Avalon
|
from avalon_framework import Avalon
|
||||||
import GPUtil
|
|
||||||
import psutil
|
|
||||||
|
|
||||||
|
|
||||||
VERSION = '3.2.0'
|
VERSION = '3.3.0'
|
||||||
|
|
||||||
LEGAL_INFO = f'''Video2X Version: {VERSION}
|
LEGAL_INFO = f'''Video2X Version: {VERSION}
|
||||||
Author: K4YT3X
|
Author: K4YT3X
|
||||||
@ -89,11 +85,6 @@ LOGO = r'''
|
|||||||
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
|
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
|
||||||
'''
|
'''
|
||||||
|
|
||||||
# each process might take up to 2.5 GB during initialization.
|
|
||||||
# (system memory, not to be confused with GPU memory)
|
|
||||||
SYS_MEM_PER_PROCESS = 2.5
|
|
||||||
GPU_MEM_PER_PROCESS = 3.5
|
|
||||||
|
|
||||||
|
|
||||||
def parse_arguments():
|
def parse_arguments():
|
||||||
"""Processes CLI arguments
|
"""Processes CLI arguments
|
||||||
@ -106,12 +97,11 @@ def parse_arguments():
|
|||||||
|
|
||||||
# video options
|
# video options
|
||||||
file_options = parser.add_argument_group('File Options')
|
file_options = parser.add_argument_group('File Options')
|
||||||
file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store')
|
file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store', required=True)
|
||||||
file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store')
|
file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store', required=True)
|
||||||
|
|
||||||
# upscaler options
|
# upscaler options
|
||||||
upscaler_options = parser.add_argument_group('Upscaler Options')
|
upscaler_options = parser.add_argument_group('Upscaler Options')
|
||||||
upscaler_options.add_argument('-m', '--method', help='upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'])
|
|
||||||
upscaler_options.add_argument('-d', '--driver', help='upscaling driver', action='store', default='waifu2x_caffe', choices=AVAILABLE_DRIVERS)
|
upscaler_options.add_argument('-d', '--driver', help='upscaling driver', action='store', default='waifu2x_caffe', choices=AVAILABLE_DRIVERS)
|
||||||
upscaler_options.add_argument('-y', '--model_dir', type=pathlib.Path, help='directory containing model JSON files', action='store')
|
upscaler_options.add_argument('-y', '--model_dir', type=pathlib.Path, help='directory containing model JSON files', action='store')
|
||||||
upscaler_options.add_argument('-p', '--processes', help='number of processes to use for upscaling', action='store', type=int, default=1)
|
upscaler_options.add_argument('-p', '--processes', help='number of processes to use for upscaling', action='store', type=int, default=1)
|
||||||
@ -139,61 +129,6 @@ def print_logo():
|
|||||||
print(f'\n{Avalon.FM.BD}{f"Version {VERSION}".rjust(36, " ")}{Avalon.FM.RST}\n')
|
print(f'\n{Avalon.FM.BD}{f"Version {VERSION}".rjust(36, " ")}{Avalon.FM.RST}\n')
|
||||||
|
|
||||||
|
|
||||||
def check_memory():
|
|
||||||
""" Check usable system memory
|
|
||||||
Warn the user if insufficient memory is available for
|
|
||||||
the number of processes that the user have chosen.
|
|
||||||
"""
|
|
||||||
|
|
||||||
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 in ['gpu', 'cudnn']:
|
|
||||||
if not (shutil.which('nvidia-smi') or
|
|
||||||
pathlib.Path(r'C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe').is_file()):
|
|
||||||
# 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 processes you\'re using')
|
|
||||||
else:
|
|
||||||
with contextlib.suppress(ValueError):
|
|
||||||
# "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_process = SYS_MEM_PER_PROCESS
|
|
||||||
else:
|
|
||||||
mem_per_process = GPU_MEM_PER_PROCESS
|
|
||||||
|
|
||||||
# if user doesn't even have enough memory to run even one process
|
|
||||||
if memory_available < mem_per_process:
|
|
||||||
Avalon.warning(f'You might have insufficient amount of {memory_type} memory available to run this program ({memory_available} GB)')
|
|
||||||
Avalon.warning('Proceed with caution')
|
|
||||||
if args.processes > 1:
|
|
||||||
if Avalon.ask('Reduce number of processes to avoid crashing?', default=True, batch=args.batch):
|
|
||||||
args.processes = 1
|
|
||||||
# if memory available is less than needed, warn the user
|
|
||||||
elif memory_available < (mem_per_process * args.processes):
|
|
||||||
Avalon.warning(f'Each waifu2x-caffe process will require up to {SYS_MEM_PER_PROCESS} GB of system memory')
|
|
||||||
Avalon.warning(f'You demanded {args.processes} processes to be created, but you only have {round(memory_available, 4)} GB {memory_type} memory available')
|
|
||||||
Avalon.warning(f'{mem_per_process * args.processes} GB of {memory_type} memory is recommended for {args.processes} processes')
|
|
||||||
Avalon.warning(f'With your current amount of {memory_type} memory available, {int(memory_available // mem_per_process)} processes is recommended')
|
|
||||||
|
|
||||||
# ask the user if he / she wants to change to the recommended
|
|
||||||
# number of processes
|
|
||||||
if Avalon.ask('Change to the recommended value?', default=True, batch=args.batch):
|
|
||||||
args.processes = int(memory_available // mem_per_process)
|
|
||||||
else:
|
|
||||||
Avalon.warning('Proceed with caution')
|
|
||||||
|
|
||||||
|
|
||||||
def read_config(config_file: pathlib.Path) -> dict:
|
def read_config(config_file: pathlib.Path) -> dict:
|
||||||
""" read video2x configurations from config file
|
""" read video2x configurations from config file
|
||||||
|
|
||||||
@ -208,36 +143,6 @@ def read_config(config_file: pathlib.Path) -> dict:
|
|||||||
return yaml.load(config, Loader=yaml.FullLoader)
|
return yaml.load(config, Loader=yaml.FullLoader)
|
||||||
|
|
||||||
|
|
||||||
def absolutify_paths(config):
|
|
||||||
""" Check to see if paths to binaries are absolute
|
|
||||||
|
|
||||||
This function checks if paths to binary files are absolute.
|
|
||||||
If not, then absolutify the path.
|
|
||||||
|
|
||||||
Arguments:
|
|
||||||
config {dict} -- configuration file dictionary
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
dict -- configuration file dictionary
|
|
||||||
"""
|
|
||||||
current_directory = pathlib.Path(sys.argv[0]).parent.absolute()
|
|
||||||
|
|
||||||
directories = [
|
|
||||||
config['waifu2x_caffe']['waifu2x_caffe_path'],
|
|
||||||
config['waifu2x_converter']['waifu2x_converter_path'],
|
|
||||||
config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'],
|
|
||||||
config['anime4k']['anime4k_path'],
|
|
||||||
config['ffmpeg']['ffmpeg_path'],
|
|
||||||
config['video2x']['video2x_cache_directory']
|
|
||||||
]
|
|
||||||
|
|
||||||
for directory in directories:
|
|
||||||
if directory and re.match('^[a-z]:', directory, re.IGNORECASE) is None:
|
|
||||||
directory = current_directory / directory
|
|
||||||
|
|
||||||
return config
|
|
||||||
|
|
||||||
|
|
||||||
# /////////////////// Execution /////////////////// #
|
# /////////////////// Execution /////////////////// #
|
||||||
|
|
||||||
# this is not a library
|
# this is not a library
|
||||||
@ -256,53 +161,9 @@ if args.version:
|
|||||||
print(LEGAL_INFO)
|
print(LEGAL_INFO)
|
||||||
sys.exit(0)
|
sys.exit(0)
|
||||||
|
|
||||||
# arguments sanity check
|
|
||||||
if not args.input:
|
|
||||||
Avalon.error('You must specify input video file/directory path')
|
|
||||||
raise ArgumentError('input video path not specified')
|
|
||||||
if not args.output:
|
|
||||||
Avalon.error('You must specify output video file/directory path')
|
|
||||||
raise ArgumentError('output video path not specified')
|
|
||||||
if (args.driver in ['waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k']) and args.width and args.height:
|
|
||||||
Avalon.error('Selected driver accepts only scaling ratio')
|
|
||||||
raise ArgumentError('selected driver supports only scaling ratio')
|
|
||||||
if args.driver == 'waifu2x_ncnn_vulkan' and args.ratio is not None and (args.ratio > 2 or not args.ratio.is_integer()):
|
|
||||||
Avalon.error('Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
|
|
||||||
raise ArgumentError('scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
|
|
||||||
if args.driver == 'srmd_ncnn_vulkan' and args.ratio is not None and (args.ratio not in [2, 3, 4]):
|
|
||||||
Avalon.error('Scaling ratio must be one of 2, 3 or 4 for srmd_ncnn_vulkan')
|
|
||||||
raise ArgumentError('scaling ratio must be one of 2, 3 or 4 for srmd_ncnn_vulkan')
|
|
||||||
if (args.width or args.height) and args.ratio:
|
|
||||||
Avalon.error('You can only specify either scaling ratio or output width and height')
|
|
||||||
raise ArgumentError('both scaling ration and width/height specified')
|
|
||||||
if (args.width and not args.height) or (not args.width and args.height):
|
|
||||||
Avalon.error('You must specify both width and height')
|
|
||||||
raise ArgumentError('only one of width or height is specified')
|
|
||||||
|
|
||||||
# check available memory if driver is waifu2x-based
|
|
||||||
if args.driver in ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan']:
|
|
||||||
check_memory()
|
|
||||||
|
|
||||||
# anime4k runs significantly faster with more processes
|
|
||||||
if args.driver == 'anime4k' and args.processes <= 1:
|
|
||||||
Avalon.warning('Anime4K runs significantly faster with more processes')
|
|
||||||
if Avalon.ask('Use more processes of Anime4K?', default=True, batch=args.batch):
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
processes = Avalon.gets('Amount of processes to use [5]: ', default=5, batch=args.batch)
|
|
||||||
args.processes = int(processes)
|
|
||||||
break
|
|
||||||
except ValueError:
|
|
||||||
if processes == '':
|
|
||||||
args.processes = 5
|
|
||||||
break
|
|
||||||
Avalon.error(f'{processes} is not a valid integer')
|
|
||||||
|
|
||||||
# read configurations from configuration file
|
# read configurations from configuration file
|
||||||
config = read_config(args.config)
|
config = read_config(args.config)
|
||||||
|
|
||||||
# config = absolutify_paths(config)
|
|
||||||
|
|
||||||
# load waifu2x configuration
|
# load waifu2x configuration
|
||||||
driver_settings = config[args.driver]
|
driver_settings = config[args.driver]
|
||||||
|
|
||||||
@ -313,34 +174,6 @@ if not pathlib.Path(driver_settings['path']).exists():
|
|||||||
Avalon.error('Please check the configuration file settings')
|
Avalon.error('Please check the configuration file settings')
|
||||||
raise FileNotFoundError(driver_settings['path'])
|
raise FileNotFoundError(driver_settings['path'])
|
||||||
|
|
||||||
# if the driver is Anime4K, check if JDK 12 is installed
|
|
||||||
jdk_available = True
|
|
||||||
if args.driver == 'anime4k':
|
|
||||||
|
|
||||||
# if specified JDK path doesn't exist
|
|
||||||
if not pathlib.Path(driver_settings['java_path']).is_file():
|
|
||||||
|
|
||||||
# try to find JDK on system
|
|
||||||
if shutil.which('java') is not None:
|
|
||||||
|
|
||||||
# check if JDK has master version 12
|
|
||||||
java_version_output = subprocess.run(['java', '-version'], capture_output=True).stderr
|
|
||||||
if re.search(r'java version "12\.\d\.\d"', java_version_output.decode().split('\n')[0]) is not None:
|
|
||||||
driver_settings['java_path'] = shutil.which('java')
|
|
||||||
else:
|
|
||||||
jdk_available = False
|
|
||||||
|
|
||||||
# if java is not found in PATH
|
|
||||||
else:
|
|
||||||
jdk_available = False
|
|
||||||
|
|
||||||
# if JDK 12 is not found
|
|
||||||
# warn the user and exit
|
|
||||||
if jdk_available is False:
|
|
||||||
Avalon.error('Cannot find JDK 12 on this system')
|
|
||||||
Avalon.error('Please ensure you have JDK 12 installed and configured')
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
# read FFmpeg configuration
|
# read FFmpeg configuration
|
||||||
ffmpeg_settings = config['ffmpeg']
|
ffmpeg_settings = config['ffmpeg']
|
||||||
|
|
||||||
@ -358,25 +191,18 @@ if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
|
|||||||
Avalon.error('Specified cache directory is a file/link')
|
Avalon.error('Specified cache directory is a file/link')
|
||||||
raise FileExistsError('Specified cache directory is a file/link')
|
raise FileExistsError('Specified cache directory is a file/link')
|
||||||
|
|
||||||
|
# if cache directory doesn't exist
|
||||||
|
# ask the user if it should be created
|
||||||
elif not video2x_cache_directory.exists():
|
elif not video2x_cache_directory.exists():
|
||||||
# if destination file is a file or a symbolic link
|
|
||||||
Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist')
|
|
||||||
|
|
||||||
# try creating the cache directory
|
|
||||||
if Avalon.ask('Create directory?', default=True, batch=args.batch):
|
|
||||||
try:
|
try:
|
||||||
|
Avalon.debug_info(f'Creating cache directory {video2x_cache_directory}')
|
||||||
video2x_cache_directory.mkdir(parents=True, exist_ok=True)
|
video2x_cache_directory.mkdir(parents=True, exist_ok=True)
|
||||||
Avalon.info(f'{video2x_cache_directory} created')
|
|
||||||
|
|
||||||
# there can be a number of exceptions here
|
# there can be a number of exceptions here
|
||||||
# PermissionError, FileExistsError, etc.
|
# PermissionError, FileExistsError, etc.
|
||||||
# therefore, we put a catch-them-all here
|
# therefore, we put a catch-them-all here
|
||||||
except Exception as e:
|
except Exception as exception:
|
||||||
Avalon.error(f'Unable to create {video2x_cache_directory}')
|
Avalon.error(f'Unable to create {video2x_cache_directory}')
|
||||||
Avalon.error('Aborting...')
|
raise exception
|
||||||
raise e
|
|
||||||
else:
|
|
||||||
raise FileNotFoundError('Could not create cache directory')
|
|
||||||
|
|
||||||
|
|
||||||
# start execution
|
# start execution
|
||||||
@ -400,10 +226,13 @@ try:
|
|||||||
Avalon.error('Suffix must be specified for FFmpeg')
|
Avalon.error('Suffix must be specified for FFmpeg')
|
||||||
raise Exception('No suffix specified')
|
raise Exception('No suffix specified')
|
||||||
|
|
||||||
upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, driver_settings=driver_settings, ffmpeg_settings=ffmpeg_settings)
|
upscaler = Upscaler(input_video=args.input,
|
||||||
|
output_video=args.output,
|
||||||
|
driver_settings=driver_settings,
|
||||||
|
ffmpeg_settings=ffmpeg_settings)
|
||||||
|
|
||||||
# set optional options
|
# set optional options
|
||||||
upscaler.waifu2x_driver = args.driver
|
upscaler.driver = args.driver
|
||||||
upscaler.scale_width = args.width
|
upscaler.scale_width = args.width
|
||||||
upscaler.scale_height = args.height
|
upscaler.scale_height = args.height
|
||||||
upscaler.scale_ratio = args.ratio
|
upscaler.scale_ratio = args.ratio
|
||||||
@ -414,9 +243,7 @@ try:
|
|||||||
upscaler.preserve_frames = preserve_frames
|
upscaler.preserve_frames = preserve_frames
|
||||||
|
|
||||||
# run upscaler
|
# run upscaler
|
||||||
upscaler.create_temp_directories()
|
|
||||||
upscaler.run()
|
upscaler.run()
|
||||||
upscaler.cleanup_temp_directories()
|
|
||||||
|
|
||||||
# if input specified is a directory
|
# if input specified is a directory
|
||||||
elif args.input.is_dir():
|
elif args.input.is_dir():
|
||||||
@ -428,10 +255,13 @@ try:
|
|||||||
|
|
||||||
for input_video in [f for f in args.input.iterdir() if f.is_file()]:
|
for input_video in [f for f in args.input.iterdir() if f.is_file()]:
|
||||||
output_video = args.output / input_video.name
|
output_video = args.output / input_video.name
|
||||||
upscaler = Upscaler(input_video=input_video, output_video=output_video, method=args.method, driver_settings=driver_settings, ffmpeg_settings=ffmpeg_settings)
|
upscaler = Upscaler(input_video=input_video,
|
||||||
|
output_video=output_video,
|
||||||
|
driver_settings=driver_settings,
|
||||||
|
ffmpeg_settings=ffmpeg_settings)
|
||||||
|
|
||||||
# set optional options
|
# set optional options
|
||||||
upscaler.waifu2x_driver = args.driver
|
upscaler.driver = args.driver
|
||||||
upscaler.scale_width = args.width
|
upscaler.scale_width = args.width
|
||||||
upscaler.scale_height = args.height
|
upscaler.scale_height = args.height
|
||||||
upscaler.scale_ratio = args.ratio
|
upscaler.scale_ratio = args.ratio
|
||||||
@ -442,9 +272,7 @@ try:
|
|||||||
upscaler.preserve_frames = preserve_frames
|
upscaler.preserve_frames = preserve_frames
|
||||||
|
|
||||||
# run upscaler
|
# run upscaler
|
||||||
upscaler.create_temp_directories()
|
|
||||||
upscaler.run()
|
upscaler.run()
|
||||||
upscaler.cleanup_temp_directories()
|
|
||||||
else:
|
else:
|
||||||
Avalon.error('Input path is neither a file nor a directory')
|
Avalon.error('Input path is neither a file nor a directory')
|
||||||
raise FileNotFoundError(f'{args.input} is neither file nor directory')
|
raise FileNotFoundError(f'{args.input} is neither file nor directory')
|
||||||
|
@ -19,8 +19,8 @@ waifu2x_caffe:
|
|||||||
model_dir: null
|
model_dir: null
|
||||||
crop_w: null
|
crop_w: null
|
||||||
crop_h: null
|
crop_h: null
|
||||||
waifu2x_converter:
|
waifu2x_converter_cpp:
|
||||||
path: 'C:\Users\K4YT3X\AppData\Local\video2x\waifu2x-converter-cpp'
|
path: 'C:\Users\K4YT3X\AppData\Local\video2x\waifu2x-converter-cpp\waifu2x-converter-cpp'
|
||||||
output-format: null
|
output-format: null
|
||||||
png-compression: null
|
png-compression: null
|
||||||
image-quality: null
|
image-quality: null
|
||||||
@ -47,9 +47,6 @@ waifu2x_ncnn_vulkan:
|
|||||||
m: models-cunet
|
m: models-cunet
|
||||||
g: 0
|
g: 0
|
||||||
j: '1:2:2'
|
j: '1:2:2'
|
||||||
anime4k:
|
|
||||||
path: 'C:\Users\K4YT3X\AppData\Local\video2x\anime4k\Anime4K.jar'
|
|
||||||
java_path: 'C:\Program Files\Java\jdk-12.0.2\bin\java.exe'
|
|
||||||
srmd_ncnn_vulkan:
|
srmd_ncnn_vulkan:
|
||||||
path: 'C:\Users\K4YT3X\AppData\Local\video2x\srmd-ncnn-vulkan\srmd-ncnn-vulkan'
|
path: 'C:\Users\K4YT3X\AppData\Local\video2x\srmd-ncnn-vulkan\srmd-ncnn-vulkan'
|
||||||
v: null
|
v: null
|
||||||
@ -62,6 +59,27 @@ srmd_ncnn_vulkan:
|
|||||||
g: 0
|
g: 0
|
||||||
j: '1:2:2'
|
j: '1:2:2'
|
||||||
x: null
|
x: null
|
||||||
|
anime4kcpp:
|
||||||
|
path: 'C:\Users\k4yt3x\AppData\Local\video2x\anime4kcpp\Anime4KCPP_CLI'
|
||||||
|
input: null
|
||||||
|
output: null
|
||||||
|
passes: 2
|
||||||
|
pushColorCount: 2
|
||||||
|
strengthColor: 0.3
|
||||||
|
strengthGradient: 1
|
||||||
|
zoomFactor: 2
|
||||||
|
threads: 16
|
||||||
|
fastMode: null
|
||||||
|
videoMode: true
|
||||||
|
preview: null
|
||||||
|
preProcessing: null
|
||||||
|
postProcessing: null
|
||||||
|
preFilters: 4
|
||||||
|
postFilters: 40
|
||||||
|
GPUMode: null
|
||||||
|
listGPUs: null
|
||||||
|
platformID: 0
|
||||||
|
deviceID: 0
|
||||||
ffmpeg:
|
ffmpeg:
|
||||||
ffmpeg_path: 'C:\Users\K4YT3X\AppData\Local\video2x\ffmpeg-latest-win64-static\bin'
|
ffmpeg_path: 'C:\Users\K4YT3X\AppData\Local\video2x\ffmpeg-latest-win64-static\bin'
|
||||||
video_to_frames:
|
video_to_frames:
|
||||||
|
Loading…
Reference in New Issue
Block a user