mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-11-04 06:31:00 +01:00
relocated files
This commit is contained in:
K4YT3X
parent
592552d82f
commit
1d9ea60e4d
@ -1,55 +0,0 @@
|
|||||||
#!/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
ffmpeg.py
89
ffmpeg.py
@ -1,89 +0,0 @@
|
|||||||
#!/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)
|
|
||||||
BIN
testvid.mp4
BIN
testvid.mp4
Binary file not shown.
@ -1,6 +0,0 @@
|
|||||||
{
|
|
||||||
"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"
|
|
||||||
}
|
|
||||||
361
video2x.py
361
video2x.py
@ -1,361 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# -*- coding: utf-8 -*-
|
|
||||||
"""
|
|
||||||
|
|
||||||
__ __ _ _ ___ __ __
|
|
||||||
\ \ / / (_) | | |__ \ \ \ / /
|
|
||||||
\ \ / / _ __| | ___ ___ ) | \ V /
|
|
||||||
\ \/ / | | / _` | / _ \ / _ \ / / > <
|
|
||||||
\ / | | | (_| | | __/ | (_) | / /_ / . \
|
|
||||||
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
|
|
||||||
|
|
||||||
|
|
||||||
Name: Video2x Controller
|
|
||||||
Author: K4YT3X
|
|
||||||
Date Created: Feb 24, 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
|
|
||||||
|
|
||||||
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 ffmpeg import Ffmpeg
|
|
||||||
from fractions import Fraction
|
|
||||||
from waifu2x import Waifu2x
|
|
||||||
import argparse
|
|
||||||
import inspect
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import psutil
|
|
||||||
import shutil
|
|
||||||
import subprocess
|
|
||||||
import threading
|
|
||||||
import time
|
|
||||||
import traceback
|
|
||||||
|
|
||||||
VERSION = '2.1.6'
|
|
||||||
|
|
||||||
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
|
|
||||||
options_group = parser.add_argument_group('Options')
|
|
||||||
options_group.add_argument('--width', help='Output video width', action='store', type=int, default=False)
|
|
||||||
options_group.add_argument('--height', help='Output video height', action='store', type=int, default=False)
|
|
||||||
options_group.add_argument('-f', '--factor', help='Factor to upscale the videos by', action='store', type=int, default=False)
|
|
||||||
options_group.add_argument('-v', '--video', help='Specify source video file', action='store', default=False)
|
|
||||||
options_group.add_argument('-o', '--output', help='Specify output file', action='store', default=False)
|
|
||||||
options_group.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb')
|
|
||||||
options_group.add_argument('-t', '--threads', help='Specify number of threads to use for upscaling', action='store', type=int, default=5)
|
|
||||||
options_group.add_argument('-c', '--config', help='Manually specify config file', action='store', default='video2x.json')
|
|
||||||
|
|
||||||
# Render drivers, at least one option must be specified
|
|
||||||
driver_group = parser.add_argument_group('Render Drivers')
|
|
||||||
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)
|
|
||||||
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():
|
|
||||||
""" Reads configuration file
|
|
||||||
|
|
||||||
Returns a dictionary read by json.
|
|
||||||
"""
|
|
||||||
with open(args.config, 'r') as raw_config:
|
|
||||||
config = json.load(raw_config)
|
|
||||||
return config
|
|
||||||
|
|
||||||
|
|
||||||
def get_video_info():
|
|
||||||
"""Gets original video information
|
|
||||||
|
|
||||||
Retrieves original video information using
|
|
||||||
ffprobe, then export it into json file.
|
|
||||||
Finally it reads, parses the json file and
|
|
||||||
returns a dictionary
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
dictionary -- original video information
|
|
||||||
"""
|
|
||||||
json_str = subprocess.check_output(
|
|
||||||
'\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(ffmpeg_path, args.video))
|
|
||||||
return json.loads(json_str.decode('utf-8'))
|
|
||||||
|
|
||||||
|
|
||||||
def check_model_type(args):
|
|
||||||
"""
|
|
||||||
Check if the model demanded from cli
|
|
||||||
argument is legal.
|
|
||||||
"""
|
|
||||||
models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
|
|
||||||
'anime_style_art_rgb', 'photo', 'anime_style_art_y']
|
|
||||||
if args.model_type not in models_available:
|
|
||||||
Avalon.error('Specified model type not found!')
|
|
||||||
Avalon.info('Available models:')
|
|
||||||
for model in models_available:
|
|
||||||
print(model)
|
|
||||||
exit(1)
|
|
||||||
|
|
||||||
|
|
||||||
def upscale_frames(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(FRAMES, f) for f in os.listdir(FRAMES) if os.path.isfile(os.path.join(FRAMES, f))]
|
|
||||||
|
|
||||||
# If we have less images than threads,
|
|
||||||
# create only the threads necessary
|
|
||||||
if len(frames) < args.threads:
|
|
||||||
args.threads = len(frames)
|
|
||||||
|
|
||||||
# Create a folder for each thread and append folder
|
|
||||||
# name into a list
|
|
||||||
|
|
||||||
thread_pool = []
|
|
||||||
for thread_id in range(args.threads):
|
|
||||||
thread_folder = '{}\\{}'.format(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], UPSCALED, args.width, args.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 video2x():
|
|
||||||
"""Main controller for Video2X
|
|
||||||
|
|
||||||
This function controls the flow of video conversion
|
|
||||||
and handles all necessary functions.
|
|
||||||
"""
|
|
||||||
|
|
||||||
check_model_type(args)
|
|
||||||
|
|
||||||
# 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))
|
|
||||||
|
|
||||||
# Initialize objects for ffmpeg and waifu2x-caffe
|
|
||||||
fm = Ffmpeg(ffmpeg_path, args.output, ffmpeg_arguments)
|
|
||||||
w2 = Waifu2x(waifu2x_path, method, args.model_type)
|
|
||||||
|
|
||||||
# Clear and create directories
|
|
||||||
if os.path.isdir(FRAMES):
|
|
||||||
shutil.rmtree(FRAMES)
|
|
||||||
if os.path.isdir(UPSCALED):
|
|
||||||
shutil.rmtree(UPSCALED)
|
|
||||||
os.mkdir(FRAMES)
|
|
||||||
os.mkdir(UPSCALED)
|
|
||||||
|
|
||||||
# Extract frames from video
|
|
||||||
fm.extract_frames(args.video, FRAMES)
|
|
||||||
|
|
||||||
Avalon.info('Reading video information')
|
|
||||||
info = 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')
|
|
||||||
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 args.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(args.factor * coded_width, args.factor * coded_height), UPSCALED)
|
|
||||||
|
|
||||||
# Use user defined output size
|
|
||||||
else:
|
|
||||||
fm.convert_video(framerate, '{}x{}'.format(args.width, args.height), UPSCALED)
|
|
||||||
Avalon.info('Conversion completed')
|
|
||||||
|
|
||||||
# Extract and press audio in
|
|
||||||
Avalon.info('Stripping audio track from original video')
|
|
||||||
fm.extract_audio(args.video, UPSCALED)
|
|
||||||
Avalon.info('Inserting audio track into new video')
|
|
||||||
fm.insert_audio_track(UPSCALED)
|
|
||||||
|
|
||||||
|
|
||||||
# /////////////////// Execution /////////////////// #
|
|
||||||
|
|
||||||
# This is not a library
|
|
||||||
if __name__ != '__main__':
|
|
||||||
Avalon.error('This file cannot be imported')
|
|
||||||
exit(1)
|
|
||||||
|
|
||||||
# Process cli arguments
|
|
||||||
args = process_arguments()
|
|
||||||
|
|
||||||
# Print video2x banner
|
|
||||||
print_logo()
|
|
||||||
|
|
||||||
# Check if arguments are valid / all necessary argument
|
|
||||||
# values are specified
|
|
||||||
if not args.video:
|
|
||||||
Avalon.error('You need to specify the video to process')
|
|
||||||
exit(1)
|
|
||||||
elif (not args.width or not args.height) and not args.factor:
|
|
||||||
Avalon.error('You must specify output video width and height or upscale factor')
|
|
||||||
exit(1)
|
|
||||||
elif not args.output:
|
|
||||||
Avalon.error('You need to specify the output video name')
|
|
||||||
exit(1)
|
|
||||||
elif not args.cpu and not args.gpu and not args.cudnn:
|
|
||||||
Avalon.error('You need to specify the enlarging processing unit')
|
|
||||||
exit(1)
|
|
||||||
|
|
||||||
# Convert paths to absolute paths
|
|
||||||
args.video = os.path.abspath(args.video)
|
|
||||||
args.output = os.path.abspath(args.output)
|
|
||||||
|
|
||||||
# Read configurations from config file
|
|
||||||
config = read_config()
|
|
||||||
waifu2x_path = config['waifu2x_path']
|
|
||||||
ffmpeg_path = config['ffmpeg_path']
|
|
||||||
ffmpeg_arguments = config['ffmpeg_arguments']
|
|
||||||
ffmpeg_hwaccel = config['ffmpeg_hwaccel']
|
|
||||||
|
|
||||||
# Add a forward slash to directory if not present
|
|
||||||
# otherwise there will be a format error
|
|
||||||
if ffmpeg_path[-1] != '/' and ffmpeg_path[-1] != '\\':
|
|
||||||
ffmpeg_path = '{}/'.format(ffmpeg_path)
|
|
||||||
|
|
||||||
# Check if FFMPEG and waifu2x are present
|
|
||||||
if not os.path.isdir(ffmpeg_path):
|
|
||||||
Avalon.error('FFMPEG binaries not found')
|
|
||||||
Avalon.error('Please specify FFMPEG binaries location in config file')
|
|
||||||
Avalon.error('Current value: {}'.format(ffmpeg_path))
|
|
||||||
raise FileNotFoundError(ffmpeg_path)
|
|
||||||
if not os.path.isfile(waifu2x_path):
|
|
||||||
Avalon.error('Waifu2x CUI executable not found')
|
|
||||||
Avalon.error('Please specify Waifu2x CUI location in config file')
|
|
||||||
Avalon.error('Current value: {}'.format(waifu2x_path))
|
|
||||||
raise FileNotFoundError(waifu2x_path)
|
|
||||||
|
|
||||||
# Check usable memory
|
|
||||||
# Warn the user if insufficient memory is available for
|
|
||||||
# the number of threads that the user have chosen. Each
|
|
||||||
# thread might take up to 2.5 GB during initialization.
|
|
||||||
MEM_PER_THREAD = 2.5
|
|
||||||
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')
|
|
||||||
|
|
||||||
# Start execution
|
|
||||||
try:
|
|
||||||
begin_time = time.time()
|
|
||||||
video2x()
|
|
||||||
Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5)))
|
|
||||||
except Exception:
|
|
||||||
Avalon.error('An exception occurred')
|
|
||||||
traceback.print_exc()
|
|
||||||
167
video2x_setup.py
167
video2x_setup.py
@ -1,167 +0,0 @@
|
|||||||
#!/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
waifu2x.py
57
waifu2x.py
@ -1,57 +0,0 @@
|
|||||||
#!/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…
x
Reference in New Issue
Block a user