mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 15:48:13 +00:00
relocated files
This commit is contained in:
K4YT3X
parent
d6f1bae2e6
commit
592552d82f
55
bin/config_generator.py
Executable file
55
bin/config_generator.py
Executable file
@ -0,0 +1,55 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: Video2x Config Generator
|
||||
Author: K4YT3X
|
||||
Date Created: October 23, 2018
|
||||
Last Modified: November 26, 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) 2018 K4YT3X
|
||||
"""
|
||||
from avalon_framework import Avalon
|
||||
import json
|
||||
import os
|
||||
|
||||
VERSION = '1.0.0'
|
||||
|
||||
|
||||
def get_path(text):
|
||||
""" Get path and validate
|
||||
"""
|
||||
while True:
|
||||
path = Avalon.gets(text)
|
||||
if os.path.isdir(path):
|
||||
return path
|
||||
Avalon.error('{} id not a directory / folder'.format(path))
|
||||
|
||||
|
||||
def enroll_settings():
|
||||
settings = {}
|
||||
|
||||
settings['waifu2x_path'] = get_path('waifu2x-caffe-cui.exe path: ')
|
||||
settings['ffmpeg_path'] = get_path('ffmpeg binaries directory: ')
|
||||
settings['ffmpeg_arguments'] = Avalon.gets('Extra arguments passed to ffmpeg: ')
|
||||
|
||||
settings['ffmpeg_hwaccel'] = Avalon.gets('ffmpeg hardware acceleration method (cuda): ')
|
||||
if settings['ffmpeg_hwaccel'] == '':
|
||||
settings['ffmpeg_hwaccel'] = 'cuda'
|
||||
|
||||
return settings
|
||||
|
||||
|
||||
def write_config(settings):
|
||||
with open('video2x.json', 'w') as config:
|
||||
json.dump(settings, config, indent=2)
|
||||
config.close()
|
||||
|
||||
|
||||
try:
|
||||
print('Video2X Config Generator {}'.format(VERSION))
|
||||
write_config(enroll_settings())
|
||||
except KeyboardInterrupt:
|
||||
Avalon.warning('Exiting...')
|
||||
89
bin/ffmpeg.py
Executable file
89
bin/ffmpeg.py
Executable file
@ -0,0 +1,89 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: FFMPEG Class
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: November 2, 2018
|
||||
|
||||
Description: This class handles all FFMPEG related
|
||||
operations.
|
||||
|
||||
Version 2.0.5
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
|
||||
class Ffmpeg:
|
||||
"""This class communicates with ffmpeg
|
||||
|
||||
This class deals with ffmpeg. It handles extracitng
|
||||
frames, stripping audio, converting images into videos
|
||||
and inserting audio tracks to videos.
|
||||
"""
|
||||
|
||||
def __init__(self, ffmpeg_path, outfile, ffmpeg_arguments, hardware_acc=False):
|
||||
self.ffmpeg_path = '{}ffmpeg.exe'.format(ffmpeg_path)
|
||||
self.outfile = outfile
|
||||
self.hardware_acc = hardware_acc
|
||||
self.ffmpeg_arguments = ffmpeg_arguments
|
||||
|
||||
def extract_frames(self, videoin, outpath):
|
||||
"""Extract every frame from original videos
|
||||
|
||||
This method extracts every frame from videoin
|
||||
using ffmpeg
|
||||
|
||||
Arguments:
|
||||
videoin {string} -- input video path
|
||||
outpath {string} -- video output folder
|
||||
"""
|
||||
execute = '\"{}\" -i \"{}\" {}\\extracted_%0d.png -y {}'.format(
|
||||
self.ffmpeg_path, videoin, outpath, ' '.join(self.ffmpeg_arguments))
|
||||
print(execute)
|
||||
subprocess.call(execute)
|
||||
|
||||
def extract_audio(self, videoin, outpath):
|
||||
"""Strips audio tracks from videos
|
||||
|
||||
This method strips audio tracks from videos
|
||||
into the output folder in aac format.
|
||||
|
||||
Arguments:
|
||||
videoin {string} -- input video path
|
||||
outpath {string} -- video output folder
|
||||
"""
|
||||
execute = '\"{}\" -i \"{}\" -vn -acodec copy {}\\output-audio.aac -y {}'.format(
|
||||
self.ffmpeg_path, videoin, outpath, ' '.join(self.ffmpeg_arguments))
|
||||
print(execute)
|
||||
subprocess.call(execute)
|
||||
|
||||
def convert_video(self, framerate, resolution, upscaled, ):
|
||||
"""Converts images into videos
|
||||
|
||||
This method converts a set of images into a
|
||||
video.
|
||||
|
||||
Arguments:
|
||||
framerate {float} -- target video framerate
|
||||
resolution {string} -- target video resolution
|
||||
upscaled {string} -- source images folder
|
||||
"""
|
||||
execute = '\"{}\" -r {} -f image2 -s {} -i {}\\extracted_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p {}\\no_audio.mp4 -y {}'.format(
|
||||
self.ffmpeg_path, framerate, resolution, upscaled, upscaled, ' '.join(self.ffmpeg_arguments))
|
||||
print(execute)
|
||||
subprocess.call(execute)
|
||||
|
||||
def insert_audio_track(self, upscaled):
|
||||
"""Insert audio into video
|
||||
|
||||
Inserts the WAV audio track stripped from
|
||||
the original video into final video.
|
||||
|
||||
Arguments:
|
||||
upscaled {string} -- upscaled image folder
|
||||
"""
|
||||
execute = '\"{}\" -i {}\\no_audio.mp4 -i {}\\output-audio.aac -shortest -codec copy {} -y {}'.format(
|
||||
self.ffmpeg_path, upscaled, upscaled, self.outfile, ' '.join(self.ffmpeg_arguments))
|
||||
print(execute)
|
||||
subprocess.call(execute)
|
||||
254
bin/upscaler.py
Normal file
254
bin/upscaler.py
Normal file
@ -0,0 +1,254 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: Video2X Upscaler
|
||||
Author: K4YT3X
|
||||
Date Created: December 10, 2018
|
||||
Last Modified: December 10, 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) 2018 K4YT3X
|
||||
"""
|
||||
|
||||
from avalon_framework import Avalon
|
||||
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)
|
||||
|
||||
|
||||
class Upscaler:
|
||||
|
||||
def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, factor=False, model_type='anime_style_art_rgb', threads=3):
|
||||
# Mandatory arguments
|
||||
self.input_video = input_video
|
||||
self.output_video = output_video
|
||||
self.method = method
|
||||
self.waifu2x_path = waifu2x_path
|
||||
self.ffmpeg_path = ffmpeg_path
|
||||
|
||||
# Optional arguments
|
||||
self.ffmpeg_arguments = ffmpeg_arguments
|
||||
self.ffmpeg_hwaccel = ffmpeg_hwaccel
|
||||
self.output_width = output_width
|
||||
self.output_height = output_height
|
||||
self.factor = factor
|
||||
self.model_type = model_type
|
||||
self.threads = threads
|
||||
|
||||
# Make temporary directories
|
||||
self.extracted_frames = tempfile.mkdtemp()
|
||||
self.upscaled_frames = tempfile.mkdtemp()
|
||||
|
||||
def _get_video_info(self):
|
||||
"""Gets input video information
|
||||
returns input video information using ffprobe in dictionary.
|
||||
"""
|
||||
json_str = subprocess.check_output('\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(self.ffmpeg_path, self.input_video))
|
||||
return json.loads(json_str.decode('utf-8'))
|
||||
|
||||
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)
|
||||
|
||||
def _check_arguments(self):
|
||||
# Check if arguments are valid / all necessary argument
|
||||
# values are specified
|
||||
if not self.input_video:
|
||||
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:
|
||||
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
|
||||
|
||||
This function upscales all the frames extracted
|
||||
by ffmpeg using the waifu2x-caffe binary.
|
||||
|
||||
Arguments:
|
||||
w2 {Waifu2x Object} -- initialized waifu2x object
|
||||
"""
|
||||
|
||||
# Create a container for all upscaler threads
|
||||
upscaler_threads = []
|
||||
|
||||
# List all images in the extracted frames
|
||||
frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))]
|
||||
|
||||
# If we have less images than threads,
|
||||
# create only the threads necessary
|
||||
if len(frames) < self.threads:
|
||||
self.threads = len(frames)
|
||||
|
||||
# Create a folder for each thread and append folder
|
||||
# name into a list
|
||||
|
||||
thread_pool = []
|
||||
for thread_id in range(self.threads):
|
||||
thread_folder = '{}\\{}'.format(self.extracted_frames, str(thread_id))
|
||||
|
||||
# Delete old folders and create new folders
|
||||
if os.path.isdir(thread_folder):
|
||||
shutil.rmtree(thread_folder)
|
||||
os.mkdir(thread_folder)
|
||||
|
||||
# Append folder path into list
|
||||
thread_pool.append((thread_folder, thread_id))
|
||||
|
||||
# Evenly distribute images into each folder
|
||||
# until there is none left in the folder
|
||||
for image in frames:
|
||||
# Move image
|
||||
shutil.move(image, thread_pool[0][0])
|
||||
# Rotate list
|
||||
thread_pool = thread_pool[-1:] + thread_pool[:-1]
|
||||
|
||||
# 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.name = thread_info[1]
|
||||
|
||||
# Add threads into the pool
|
||||
upscaler_threads.append(thread)
|
||||
|
||||
# Start all threads
|
||||
for thread in upscaler_threads:
|
||||
thread.start()
|
||||
|
||||
# Wait for threads to finish
|
||||
for thread in upscaler_threads:
|
||||
thread.join()
|
||||
|
||||
def run(self):
|
||||
"""Main controller for Video2X
|
||||
|
||||
This function controls the flow of video conversion
|
||||
and handles all necessary functions.
|
||||
"""
|
||||
|
||||
# Parse arguments for waifu2x
|
||||
# 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)
|
||||
self.output_video = os.path.abspath(self.output_video)
|
||||
|
||||
# Add a forward slash to directory if not present
|
||||
# otherwise there will be a format error
|
||||
if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
|
||||
self.ffmpeg_path = '{}/'.format(self.ffmpeg_path)
|
||||
|
||||
# Check if FFMPEG and waifu2x are present
|
||||
if not os.path.isdir(self.ffmpeg_path):
|
||||
raise FileNotFoundError(self.ffmpeg_path)
|
||||
if not os.path.isfile(self.waifu2x_path):
|
||||
raise FileNotFoundError(self.waifu2x_path)
|
||||
|
||||
# Initialize objects for ffmpeg and waifu2x-caffe
|
||||
fm = Ffmpeg(self.ffmpeg_path, self.video_output, self.ffmpeg_arguments)
|
||||
w2 = Waifu2x(self.waifu2x_path, self.method, self.model_type)
|
||||
|
||||
# Extract frames from video
|
||||
fm.extract_frames(self.input_video, self.extracted_frames)
|
||||
|
||||
Avalon.info('Reading video information')
|
||||
info = self._get_video_info()
|
||||
# Analyze original video with ffprobe and retrieve framerate
|
||||
# width, height = info['streams'][0]['width'], info['streams'][0]['height']
|
||||
|
||||
# Find index of video stream
|
||||
video_stream_index = None
|
||||
for stream in info['streams']:
|
||||
if stream['codec_type'] == 'video':
|
||||
video_stream_index = stream['index']
|
||||
break
|
||||
|
||||
# Exit if no video stream found
|
||||
if video_stream_index is None:
|
||||
Avalon.error('Aborting: No video stream found')
|
||||
|
||||
# Get average frame rate of video stream
|
||||
framerate = float(Fraction(info['streams'][video_stream_index]['avg_frame_rate']))
|
||||
Avalon.info('Framerate: {}'.format(framerate))
|
||||
|
||||
# Upscale images one by one using waifu2x
|
||||
Avalon.info('Starting to upscale extracted images')
|
||||
self._upscale_frames(w2)
|
||||
Avalon.info('Upscaling completed')
|
||||
|
||||
# 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)
|
||||
Avalon.info('Conversion completed')
|
||||
|
||||
# Extract and press audio in
|
||||
Avalon.info('Stripping audio track from original video')
|
||||
fm.extract_audio(self.input_video, self.upscaled_frames)
|
||||
Avalon.info('Inserting audio track into new video')
|
||||
fm.insert_audio_track(self.upscaled_frames)
|
||||
6
bin/video2x.json
Normal file
6
bin/video2x.json
Normal file
@ -0,0 +1,6 @@
|
||||
{
|
||||
"waifu2x_path": "C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe",
|
||||
"ffmpeg_path": "C:/Program Files (x86)/ffmpeg/bin/",
|
||||
"ffmpeg_arguments": [],
|
||||
"ffmpeg_hwaccel": "cuda"
|
||||
}
|
||||
140
bin/video2x.py
Executable file
140
bin/video2x.py
Executable file
@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
|
||||
__ __ _ _ ___ __ __
|
||||
\ \ / / (_) | | |__ \ \ \ / /
|
||||
\ \ / / _ __| | ___ ___ ) | \ V /
|
||||
\ \/ / | | / _` | / _ \ / _ \ / / > <
|
||||
\ / | | | (_| | | __/ | (_) | / /_ / . \
|
||||
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
|
||||
|
||||
|
||||
Name: Video2x Controller
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: December 10, 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) 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.
|
||||
"""
|
||||
from avalon_framework import Avalon
|
||||
from upscaler import Upscaler
|
||||
import argparse
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
|
||||
VERSION = '2.1.7'
|
||||
|
||||
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
|
||||
|
||||
|
||||
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')
|
||||
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('-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.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()
|
||||
|
||||
|
||||
def print_logo():
|
||||
print('__ __ _ _ ___ __ __')
|
||||
print('\\ \\ / / (_) | | |__ \\ \\ \\ / /')
|
||||
print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /')
|
||||
print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <')
|
||||
print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\')
|
||||
print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\')
|
||||
print('\n Video2X Video Enlarger')
|
||||
spaces = ((44 - len("Version " + VERSION)) // 2) * " "
|
||||
print('{}\n{} Version {}\n{}'.format(Avalon.FM.BD, spaces, VERSION, Avalon.FM.RST))
|
||||
|
||||
|
||||
def read_config(config_file):
|
||||
""" Reads configuration file
|
||||
|
||||
Returns a dictionary read by json.
|
||||
"""
|
||||
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')
|
||||
exit(1)
|
||||
|
||||
print_logo()
|
||||
|
||||
# Process cli arguments
|
||||
args = process_arguments()
|
||||
|
||||
# Read configurations from JSON
|
||||
config = read_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()
|
||||
|
||||
Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5)))
|
||||
except Exception:
|
||||
Avalon.error('An exception occurred')
|
||||
traceback.print_exc()
|
||||
167
bin/video2x_setup.py
Normal file
167
bin/video2x_setup.py
Normal file
@ -0,0 +1,167 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: Video2X Setup Script
|
||||
Author: K4YT3X
|
||||
Date Created: November 28, 2018
|
||||
Last Modified: November 29, 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) 2018 K4YT3X
|
||||
|
||||
Description: This script helps installing all dependencies of video2x
|
||||
and generates a configuration for it.
|
||||
|
||||
Installation Details:
|
||||
- ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg
|
||||
- waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe
|
||||
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import traceback
|
||||
import zipfile
|
||||
|
||||
# Requests doesn't come with windows, therefore
|
||||
# it will be installed as a dependency and imported
|
||||
# later in the script.
|
||||
# import requests
|
||||
|
||||
VERSION = '1.0.0'
|
||||
|
||||
|
||||
class Video2xSetup:
|
||||
""" Install dependencies for video2x video enlarger
|
||||
|
||||
This library is meant to be executed as a stand-alone
|
||||
script. All files will be installed under %LOCALAPPDATA%\\video2x.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.trash = []
|
||||
|
||||
def run(self):
|
||||
|
||||
print('\nInstalling Python libraries')
|
||||
self._install_python_requirements()
|
||||
|
||||
print('\nInstalling FFMPEG')
|
||||
self._install_ffmpeg()
|
||||
|
||||
print('\nInstalling waifu2x-caffe')
|
||||
self._install_waifu2x_caffe()
|
||||
|
||||
print('\nGenerating Video2X configuration file')
|
||||
self._generate_config()
|
||||
|
||||
print('\nCleaning up temporary files')
|
||||
self._cleanup()
|
||||
|
||||
def _install_python_requirements(self):
|
||||
""" Read requirements.txt and return its content
|
||||
"""
|
||||
with open('requirements.txt', 'r') as req:
|
||||
for line in req:
|
||||
package = line.split(' ')[0]
|
||||
pip_install(package)
|
||||
|
||||
def _cleanup(self):
|
||||
""" Cleanup all the temp files downloaded
|
||||
"""
|
||||
for file in self.trash:
|
||||
try:
|
||||
print('Deleting: {}'.format(file))
|
||||
os.remove(file)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def _install_ffmpeg(self):
|
||||
""" Install FFMPEG
|
||||
"""
|
||||
latest_release = 'https://ffmpeg.zeranoe.com/builds/win64/static/ffmpeg-4.1-win64-static.zip'
|
||||
|
||||
ffmpeg_zip = download(latest_release, tempfile.gettempdir())
|
||||
self.trash.append(ffmpeg_zip)
|
||||
|
||||
with zipfile.ZipFile(ffmpeg_zip) as zipf:
|
||||
zipf.extractall('{}\\video2x'.format(os.getenv('localappdata')))
|
||||
|
||||
def _install_waifu2x_caffe(self):
|
||||
""" Install waifu2x_caffe
|
||||
"""
|
||||
import requests
|
||||
|
||||
# Get latest release of waifu2x-caffe via GitHub API
|
||||
latest_release = json.loads(requests.get('https://api.github.com/repos/lltcggie/waifu2x-caffe/releases/latest').content)
|
||||
|
||||
for a in latest_release['assets']:
|
||||
if 'waifu2x-caffe.zip' in a['browser_download_url']:
|
||||
waifu2x_caffe_zip = download(a['browser_download_url'], tempfile.gettempdir())
|
||||
self.trash.append(waifu2x_caffe_zip)
|
||||
|
||||
with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
|
||||
zipf.extractall('{}\\video2x'.format(os.getenv('localappdata')))
|
||||
|
||||
def _generate_config(self):
|
||||
""" Generate video2x config
|
||||
"""
|
||||
settings = {}
|
||||
|
||||
settings['waifu2x_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(os.getenv('localappdata'))
|
||||
settings['ffmpeg_path'] = '{}\\video2x\\ffmpeg-4.1-win64-static\\bin'.format(os.getenv('localappdata'))
|
||||
settings['ffmpeg_arguments'] = ''
|
||||
settings['ffmpeg_hwaccel'] = 'cuda'
|
||||
|
||||
with open('video2x.json', 'w') as config:
|
||||
json.dump(settings, config, indent=2)
|
||||
config.close()
|
||||
|
||||
|
||||
def download(url, save_path, chunk_size=4096):
|
||||
""" Download file to local with requests library
|
||||
"""
|
||||
import requests
|
||||
|
||||
output_file = '{}\\{}'.format(save_path, url.split('/')[-1])
|
||||
print('Downloading: {}'.format(url))
|
||||
print('Chunk size: {}'.format(chunk_size))
|
||||
print('Saving to: {}'.format(output_file))
|
||||
|
||||
stream = requests.get(url, stream=True)
|
||||
|
||||
# Write content into file
|
||||
with open(output_file, 'wb') as output:
|
||||
for chunk in stream.iter_content(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
print('!', end='')
|
||||
output.write(chunk)
|
||||
print()
|
||||
|
||||
return output_file
|
||||
|
||||
|
||||
def pip_install(package):
|
||||
""" Install python package via python pip module
|
||||
|
||||
pip.main() is not available after pip 9.0.1, thus
|
||||
pip module is not used in this case.
|
||||
"""
|
||||
return subprocess.run(['python', '-m', 'pip', 'install', '-U', package]).returncode
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
print('Video2x Setup Script')
|
||||
print('Version: {}'.format(VERSION))
|
||||
setup = Video2xSetup()
|
||||
setup.run()
|
||||
print('\n Script finished successfully')
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
print('An error has occurred')
|
||||
print('Video2X Automatic Setup has failed')
|
||||
exit(1)
|
||||
57
bin/waifu2x.py
Executable file
57
bin/waifu2x.py
Executable file
@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: FFMPEG Class
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: October 23, 2018
|
||||
|
||||
Description: This class controls waifu2x
|
||||
engine
|
||||
|
||||
Version 2.0.4
|
||||
"""
|
||||
from avalon_framework import Avalon
|
||||
import subprocess
|
||||
import threading
|
||||
|
||||
|
||||
class Waifu2x:
|
||||
"""This class communicates with waifu2x cui engine
|
||||
|
||||
An object will be created for this class, containing information
|
||||
about the binary address and the processing method. When being called
|
||||
by the main program, other detailed information will be passed to
|
||||
the upscale function.
|
||||
"""
|
||||
|
||||
def __init__(self, waifu2x_path, method, model_type):
|
||||
self.waifu2x_path = waifu2x_path
|
||||
self.method = method
|
||||
self.model_type = model_type
|
||||
self.print_lock = threading.Lock()
|
||||
|
||||
def upscale(self, folderin, folderout, width, height):
|
||||
"""This is the core function for WAIFU2X class
|
||||
|
||||
Arguments:
|
||||
folderin {string} -- source folder path
|
||||
folderout {string} -- output folder path
|
||||
width {int} -- output video width
|
||||
height {int} -- output video height
|
||||
"""
|
||||
|
||||
# Print thread start message
|
||||
self.print_lock.acquire()
|
||||
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))
|
||||
self.print_lock.release()
|
||||
|
||||
# Create string for execution
|
||||
execute = '\"{}\" -p {} -I png -i \"{}\" -e png -o {} -w {} -h {} -n 3 -m noise_scale -y {}'.format(
|
||||
self.waifu2x_path, self.method, folderin, folderout, width, height, self.model_type)
|
||||
subprocess.call(execute)
|
||||
|
||||
# Print thread exiting message
|
||||
self.print_lock.acquire()
|
||||
Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name))
|
||||
self.print_lock.release()
|
||||
Loading…
Reference in New Issue
Block a user