OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2025-01-01 10:29:09 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'master' into master

Browse Source
This commit is contained in:
K4YT3X 2019-06-15 14:42:56 -04:00 committed by GitHub
commit e003cbc568
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 43 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
bin/ffmpeg.py
View File

@ -4,7 +4,7 @@
Name: FFMPEG Class Name: FFMPEG Class
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: June 5, 2019 Last Modified: June 15, 2019
Description: This class handles all FFMPEG related Description: This class handles all FFMPEG related
operations. operations.
@ -12,6 +12,7 @@ operations.
from avalon_framework import Avalon from avalon_framework import Avalon
import json import json
import subprocess import subprocess
import os
class Ffmpeg: class Ffmpeg:
@ -26,13 +27,9 @@ class Ffmpeg:
self.ffmpeg_settings = ffmpeg_settings self.ffmpeg_settings = ffmpeg_settings
self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path'] self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path']
# 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 = f'{self.ffmpeg_path}\\'
self.ffmpeg_binary = f'{self.ffmpeg_path}ffmpeg.exe' self.ffmpeg_binary = os.path.join(self.ffmpeg_path, 'ffmpeg.exe')
self.ffmpeg_probe_binary = f'{self.ffmpeg_path}ffprobe.exe' self.ffmpeg_probe_binary = os.path.join(self.ffmpeg_path, 'ffprobe.exe')
self.image_format = image_format self.image_format = image_format
def get_video_info(self, input_video): def get_video_info(self, input_video):
@ -88,7 +85,7 @@ class Ffmpeg:
execute.extend(self._read_configuration(phase='video_to_frames', section='output_options')) execute.extend(self._read_configuration(phase='video_to_frames', section='output_options'))
execute.extend([ execute.extend([
f'{extracted_frames}\\extracted_%0d.{self.image_format}' os.path.join(extracted_frames, f'extracted_%0d.{self.image_format}')
]) ])
execute.extend(self._read_configuration(phase='video_to_frames')) execute.extend(self._read_configuration(phase='video_to_frames'))
@ -120,7 +117,7 @@ class Ffmpeg:
# append input frames path into command # append input frames path into command
execute.extend([ execute.extend([
'-i', '-i',
f'{upscaled_frames}\\extracted_%d.{self.image_format}' os.path.join(upscaled_frames, f'extracted_%d.{self.image_format}')
]) ])
# read FFmpeg output options # read FFmpeg output options
@ -131,7 +128,7 @@ class Ffmpeg:
# specify output file location # specify output file location
execute.extend([ execute.extend([
f'{upscaled_frames}\\no_audio.mp4' os.path.join(upscaled_frames, 'no_audio.mp4')
]) ])
self._execute(execute) self._execute(execute)
@ -147,7 +144,7 @@ class Ffmpeg:
execute = [ execute = [
self.ffmpeg_binary, self.ffmpeg_binary,
'-i', '-i',
f'{upscaled_frames}\\no_audio.mp4', os.path.join(upscaled_frames, 'no_audio.mp4'),
'-i', '-i',
input_video input_video
] ]

10
bin/upscaler.py
View File

@ -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: June 13, 2019 Last Modified: June 15, 2019
Licensed under the GNU General Public License Version 3 (GNU GPL v3), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -53,7 +53,7 @@ class Upscaler:
self.scale_ratio = None self.scale_ratio = None
self.model_dir = None self.model_dir = None
self.threads = 5 self.threads = 5
self.video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x' self.video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
self.image_format = 'png' self.image_format = 'png'
self.preserve_frames = False self.preserve_frames = False
@ -144,7 +144,7 @@ class Upscaler:
# if this thread is not empty, then an exception has occured # if this thread is not empty, then an exception has occured
self.upscaler_exceptions = [] self.upscaler_exceptions = []
# initialize waifu2x driver # initialize waifu2x driver
if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter': if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter':
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}') raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
@ -159,7 +159,7 @@ class Upscaler:
w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions) w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions)
for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]: for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]:
renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image) renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image)
shutil.move(f'{self.upscaled_frames}\\{image}', f'{self.upscaled_frames}\\{renamed}') shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
progress_bar.join() progress_bar.join()
@ -182,7 +182,7 @@ class Upscaler:
thread_pool = [] thread_pool = []
thread_directories = [] thread_directories = []
for thread_id in range(self.threads): for thread_id in range(self.threads):
thread_directory = f'{self.extracted_frames}\\{str(thread_id)}' thread_directory = os.path.join(self.extracted_frames, str(thread_id))
thread_directories.append(thread_directory) thread_directories.append(thread_directory)
# delete old directories and create new directories # delete old directories and create new directories

22
bin/video2x.py
View File

@ -13,7 +13,10 @@ __ __ _ _ ___ __ __
Name: Video2X Controller Name: Video2X Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: June 13, 2019 Last Modified: June 15, 2019
Dev: BrianPetkovsek
Dev: SAT3LL
Licensed under the GNU General Public License Version 3 (GNU GPL v3), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -80,7 +83,7 @@ def process_arguments():
upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter']) upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter'])
upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store') upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store')
upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1) upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1)
upscaler_options.add_argument('-c', '--config', help='video2x config file location', action='store', default=f'{os.path.dirname(os.path.abspath(sys.argv[0]))}\\video2x.json') upscaler_options.add_argument('-c', '--config', help='video2x config file location', action='store', default=os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'video2x.json'))
upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true') upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true')
# scaling options # scaling options
@ -123,7 +126,8 @@ def check_memory():
# check if Nvidia-smi is available # check if Nvidia-smi is available
# GPUtil requires nvidia-smi.exe to interact with GPU # GPUtil requires nvidia-smi.exe to interact with GPU
if args.method == 'gpu' or args.method == 'cudnn': if args.method == 'gpu' or args.method == 'cudnn':
if not os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'): if not (shutil.which('nvidia-smi') or
os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe')):
# Nvidia System Management Interface not available # Nvidia System Management Interface not available
Avalon.warning('Nvidia-smi not available, skipping available memory check') Avalon.warning('Nvidia-smi not available, skipping available memory check')
Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using') Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
@ -191,20 +195,20 @@ def absolutify_paths(config):
# check waifu2x-caffe path # check waifu2x-caffe path
if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
config['waifu2x_caffe']['waifu2x_caffe_path'] = f'{current_directory}\\{config["waifu2x_caffe"]["waifu2x_caffe_path"]}' config['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(current_directory, config['waifu2x_caffe']['waifu2x_caffe_path'])
# check waifu2x-converter-cpp path # check waifu2x-converter-cpp path
if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
config['waifu2x_converter']['waifu2x_converter_path'] = f'{current_directory}\\{config["waifu2x_converter"]["waifu2x_converter_path"]}' config['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(current_directory, config['waifu2x_converter']['waifu2x_converter_path'])
# check ffmpeg path # check ffmpeg path
if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
config['ffmpeg']['ffmpeg_path'] = f'{current_directory}\\{config["ffmpeg"]["ffmpeg_path"]}' config['ffmpeg']['ffmpeg_path'] = os.path.join(current_directory, config['ffmpeg']['ffmpeg_path'])
# check video2x cache path # check video2x cache path
if config['video2x']['video2x_cache_directory']: if config['video2x']['video2x_cache_directory']:
if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE): if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
config['video2x']['video2x_cache_directory'] = f'{current_directory}\\{config["video2x"]["video2x_cache_directory"]}' config['video2x']['video2x_cache_directory'] = os.path.join(current_directory, config['video2x']['video2x_cache_directory'])
return config return config
@ -279,7 +283,7 @@ preserve_frames = config['video2x']['preserve_frames']
# create temp directories if they don't exist # create temp directories if they don't exist
if not video2x_cache_directory: if not video2x_cache_directory:
video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x' video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
if video2x_cache_directory and not os.path.isdir(video2x_cache_directory): if video2x_cache_directory and not os.path.isdir(video2x_cache_directory):
if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory): if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory):
@ -341,7 +345,7 @@ try:
# upscale videos in a directory # upscale videos in a directory
Avalon.info(f'Upscaling videos in directory: {args.input}') Avalon.info(f'Upscaling videos in directory: {args.input}')
for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]: for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]:
output_video = f'{args.output}\\{input_video}' output_video = os.path.join(args.output, input_video)
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings) upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
# set optional options # set optional options

22
bin/video2x_setup.py
View File

@ -5,7 +5,7 @@ Name: Video2X Setup Script
Author: K4YT3X Author: K4YT3X
Author: BrianPetkovsek Author: BrianPetkovsek
Date Created: November 28, 2018 Date Created: November 28, 2018
Last Modified: June 14, 2019 Last Modified: June 15, 2019
Licensed under the GNU General Public License Version 3 (GNU GPL v3), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -108,7 +108,7 @@ class Video2xSetup:
self.trash.append(ffmpeg_zip) self.trash.append(ffmpeg_zip)
with zipfile.ZipFile(ffmpeg_zip) as zipf: with zipfile.ZipFile(ffmpeg_zip) as zipf:
zipf.extractall('{}\\video2x'.format(os.getenv("localappdata"))) zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
def _install_waifu2x_caffe(self): def _install_waifu2x_caffe(self):
""" Install waifu2x_caffe """ Install waifu2x_caffe
@ -125,7 +125,7 @@ class Video2xSetup:
self.trash.append(waifu2x_caffe_zip) self.trash.append(waifu2x_caffe_zip)
with zipfile.ZipFile(waifu2x_caffe_zip) as zipf: with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
zipf.extractall('{}\\video2x'.format(os.getenv("localappdata"))) zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
def _install_waifu2x_converter_cpp(self): def _install_waifu2x_converter_cpp(self):
""" Install waifu2x_caffe """ Install waifu2x_caffe
@ -143,7 +143,7 @@ class Video2xSetup:
self.trash.append(waifu2x_converter_cpp_zip) self.trash.append(waifu2x_converter_cpp_zip)
with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf: with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
zipf.extractall('{}\\video2x\\waifu2x-converter-cpp'.format(os.getenv("localappdata"))) zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp'))
def _generate_config(self): def _generate_config(self):
""" Generate video2x config """ Generate video2x config
@ -157,14 +157,14 @@ class Video2xSetup:
# configure only the specified drivers # configure only the specified drivers
if self.driver == 'all': if self.driver == 'all':
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(local_app_data) template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
template_dict['waifu2x_converter']['waifu2x_converter_path'] = '{}\\video2x\\waifu2x-converter-cpp'.format(local_app_data) template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
elif self.driver == 'waifu2x_caffe': elif self.driver == 'waifu2x_caffe':
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(local_app_data) template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
elif self.driver == 'waifu2x_converter': elif self.driver == 'waifu2x_converter':
template_dict['waifu2x_converter']['waifu2x_converter_path'] = '{}\\video2x\\waifu2x-converter-cpp'.format(local_app_data) template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
template_dict['ffmpeg']['ffmpeg_path'] = '{}\\video2x\\ffmpeg-latest-win64-static\\bin'.format(local_app_data) template_dict['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x', 'ffmpeg-latest-win64-static', 'bin')
template_dict['video2x']['video2x_cache_directory'] = None template_dict['video2x']['video2x_cache_directory'] = None
template_dict['video2x']['preserve_frames'] = False template_dict['video2x']['preserve_frames'] = False
@ -180,7 +180,7 @@ def download(url, save_path, chunk_size=4096):
from tqdm import tqdm from tqdm import tqdm
import requests import requests
output_file = '{}\\{}'.format(save_path, url.split("/")[-1]) output_file = os.path.join(save_path, url.split('/')[-1])
print('Downloading: {}'.format(url)) print('Downloading: {}'.format(url))
print('Chunk size: {}'.format(chunk_size)) print('Chunk size: {}'.format(chunk_size))
print('Saving to: {}'.format(output_file)) print('Saving to: {}'.format(output_file))
@ -208,7 +208,7 @@ def pip_install(file):
return subprocess.run([sys.executable, '-m', 'pip', 'install', '-U', '-r', file]).returncode return subprocess.run([sys.executable, '-m', 'pip', 'install', '-U', '-r', file]).returncode
if __name__ == "__main__": if __name__ == '__main__':
try: try:
args = process_arguments() args = process_arguments()
print('Video2X Setup Script') print('Video2X Setup Script')

2
bin/waifu2x_caffe.py
View File

@ -4,7 +4,7 @@
Name: Waifu2x Caffe Driver Name: Waifu2x Caffe Driver
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: April 28, 2019 Last Modified: June 15, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-caffe. for waifu2x-caffe.

8
bin/waifu2x_converter.py
View File

@ -4,12 +4,13 @@
Name: Waifu2x Converter CPP Driver Name: Waifu2x Converter CPP Driver
Author: K4YT3X Author: K4YT3X
Date Created: February 8, 2019 Date Created: February 8, 2019
Last Modified: April 28, 2019 Last Modified: June 15, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-converter-cpp. for waifu2x-converter-cpp.
""" """
from avalon_framework import Avalon from avalon_framework import Avalon
import os
import subprocess import subprocess
import threading import threading
@ -59,7 +60,8 @@ class Waifu2xConverter:
# models_rgb must be specified manually for waifu2x-converter-cpp # models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically # if it's not specified in the arguments, create automatically
if self.waifu2x_settings['model-dir'] is None: if self.waifu2x_settings['model-dir'] is None:
self.waifu2x_settings['model-dir'] = f'{self.waifu2x_settings["waifu2x_converter_path"]}\\models_rgb' self.waifu2x_settings['model-dir'] = os.path.join(self.waifu2x_settings['waifu2x_converter_path'],
'models_rgb')
# print thread start message # print thread start message
self.print_lock.acquire() self.print_lock.acquire()
@ -75,7 +77,7 @@ class Waifu2xConverter:
# the key doesn't need to be passed in this case # the key doesn't need to be passed in this case
if key == 'waifu2x_converter_path': if key == 'waifu2x_converter_path':
execute.append(f'{str(value)}\\waifu2x-converter-cpp.exe') execute.append(os.path.join(str(value), 'waifu2x-converter-cpp.exe'))
# null or None means that leave this option out (keep default) # null or None means that leave this option out (keep default)
elif value is None or value is False: elif value is None or value is False: