Merge pull request #129 from k4yt3x/2.9.0

2.9.0
This commit is contained in:
K4YT3X 2019-07-27 21:34:35 +00:00 committed by GitHub
commit 734cf7bb0e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 265 additions and 194 deletions

View File

@ -18,10 +18,16 @@ Component names that are **bolded** can be automatically downloaded and configur
2. AMD GPU / Nvidia GPU
3. AMD GPU driver / Nvidia GPU driver / Nvidia CUDNN
4. [**FFmpeg**](https://ffmpeg.zeranoe.com/builds/)
5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases)
5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases) / [**waifu2x-ncnn-vulkan**](https://github.com/nihui/waifu2x-ncnn-vulkan)
## Recent Changes
### 2.9.0 (July 27, 2019)
- Changed file handling method from `os` to `pathlib`
- Removed f_string dependency and support for legacy versions of Python
- Organized file import statements
### 2.8.1 (July 9, 2019)
- Added automatic pixel format detection
@ -35,12 +41,6 @@ Component names that are **bolded** can be automatically downloaded and configur
- Fixed video2x custom temp folder bug found by @cr08 .
### 2.7.0 (March 30, 2019)
- Added support for different extracted image formats.
- Redesigned FFmpeg wrapper, FFmpeg settings are now customizable in the `video2x.json` config file.
- Other minor enhancements and adjustments (e.g. argument -> method variable)
### Setup Script 1.3.0 (June 25, 2019)
- Added automatic installation support for `waifu2x-ncnn-vulkan`

View File

@ -1,12 +1,10 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Video2X Exceptions
Dev: K4YT3X
Date Created: December 13, 2018
Last Modified: March 19, 2019
Last Modified: July 27, 2019
"""

View File

@ -1,25 +1,28 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: FFMPEG Class
Name: Video2X FFmpeg Controller
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: July 26, 2019
Last Modified: July 27, 2019
Description: This class handles all FFmpeg related operations.
"""
from avalon_framework import Avalon
# built-in imports
import json
import subprocess
import os
import pathlib
import subprocess
# third-party imports
from avalon_framework import Avalon
class Ffmpeg:
"""This class communicates with FFmpeg
This class deals with FFmpeg. It handles extracitng
This class deals with FFmpeg. It handles extracting
frames, stripping audio, converting images into videos
and inserting audio tracks to videos.
"""
@ -27,18 +30,17 @@ class Ffmpeg:
def __init__(self, ffmpeg_settings, image_format):
self.ffmpeg_settings = ffmpeg_settings
self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path']
self.ffmpeg_binary = os.path.join(self.ffmpeg_path, 'ffmpeg.exe')
self.ffmpeg_probe_binary = os.path.join(self.ffmpeg_path, 'ffprobe.exe')
self.ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path'])
self.ffmpeg_binary = self.ffmpeg_path / 'ffmpeg.exe'
self.ffmpeg_probe_binary = self.ffmpeg_path / 'ffprobe.exe'
self.image_format = image_format
self.pixel_format = None
def get_pixel_formats(self):
""" Get a dictionary of supported pixel formats
List all supported pixel formats and their corresponding
bit depth.
List all supported pixel formats and their
corresponding bit depth.
Returns:
dictionary -- JSON dict of all pixel formats to bit depth
@ -50,6 +52,9 @@ class Ffmpeg:
'-pix_fmts'
]
# turn elements into str
execute = [str(e) for e in execute]
Avalon.debug_info(f'Executing: {" ".join(execute)}')
# initialize dictionary to store pixel formats
@ -71,7 +76,7 @@ class Ffmpeg:
""" Gets input video information
This method reads input video information
using ffprobe in dictionary.
using ffprobe in dictionary
Arguments:
input_video {string} -- input video file path
@ -94,6 +99,9 @@ class Ffmpeg:
input_video
]
# turn elements into str
execute = [str(e) for e in execute]
Avalon.debug_info(f'Executing: {" ".join(execute)}')
json_str = subprocess.run(execute, check=True, stdout=subprocess.PIPE).stdout
return json.loads(json_str.decode('utf-8'))
@ -101,8 +109,7 @@ class Ffmpeg:
def extract_frames(self, input_video, extracted_frames):
"""Extract every frame from original videos
This method extracts every frame from videoin
using FFmpeg
This method extracts every frame from input video using FFmpeg
Arguments:
input_video {string} -- input video path
@ -120,7 +127,7 @@ class Ffmpeg:
execute.extend(self._read_configuration(phase='video_to_frames', section='output_options'))
execute.extend([
os.path.join(extracted_frames, f'extracted_%0d.{self.image_format}')
extracted_frames / f'extracted_%0d.{self.image_format}'
])
execute.extend(self._read_configuration(phase='video_to_frames'))
@ -130,8 +137,7 @@ class Ffmpeg:
def convert_video(self, framerate, resolution, upscaled_frames):
"""Converts images into videos
This method converts a set of images into a
video.
This method converts a set of images into a video
Arguments:
framerate {float} -- target video framerate
@ -150,17 +156,19 @@ class Ffmpeg:
execute.extend(self._read_configuration(phase='frames_to_video', section='input_options'))
# WORKAROUND FOR WAIFU2X-NCNN-VULKAN
# Dev: SAT3LL
# rename all .png.png suffixes to .png
import re
import shutil
regex = re.compile(r'\.png\.png$')
for raw_frame in os.listdir(upscaled_frames):
shutil.move(os.path.join(upscaled_frames, raw_frame), os.path.join(upscaled_frames, regex.sub('.png', raw_frame)))
for frame_name in upscaled_frames.iterdir():
(upscaled_frames / frame_name).rename(upscaled_frames / regex.sub('.png', str(frame_name)))
# END WORKAROUND
# append input frames path into command
execute.extend([
'-i',
os.path.join(upscaled_frames, f'extracted_%d.{self.image_format}')
upscaled_frames / f'extracted_%d.{self.image_format}'
])
# read FFmpeg output options
@ -171,7 +179,7 @@ class Ffmpeg:
# specify output file location
execute.extend([
os.path.join(upscaled_frames, 'no_audio.mp4')
upscaled_frames / 'no_audio.mp4'
])
self._execute(execute)
@ -187,7 +195,7 @@ class Ffmpeg:
execute = [
self.ffmpeg_binary,
'-i',
os.path.join(upscaled_frames, 'no_audio.mp4'),
upscaled_frames / 'no_audio.mp4',
'-i',
input_video
]
@ -271,4 +279,8 @@ class Ffmpeg:
int -- execution return code
"""
Avalon.debug_info(f'Executing: {execute}')
# turn all list elements into string to avoid errors
execute = [str(e) for e in execute]
return subprocess.run(execute, shell=True, check=True).returncode

View File

@ -1,18 +1,17 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Video2X Image Cleaner
Author: BrianPetkovsek
Author: K4YT3X
Date Created: March 24, 2019
Last Modified: April 28, 2019
Last Modified: July 27, 2019
Description: This class is to remove the extracted frames
that have already been upscaled.
"""
# built-in imports
import os
import threading
import time
@ -60,19 +59,19 @@ class ImageCleaner(threading.Thread):
"""
# list all images in the extracted frames
output_frames = [f for f in os.listdir(self.output_directory) if os.path.isfile(os.path.join(self.output_directory, f))]
output_frames = [f for f in self.output_directory.iterdir() if f.is_file()]
# compare and remove frames downscaled images that finished being upscaled
# within each thread's extracted frames directory
for i in range(self.threads):
dir_path = os.path.join(self.input_directory, str(i))
for thread_id in range(self.threads):
dir_path = self.input_directory / str(thread_id)
# for each file within all the directories
for f in os.listdir(dir_path):
file_path = os.path.join(dir_path, f)
for file in dir_path.iterdir():
file_path = dir_path / file
# if file also exists in the output directory, then the file
# has already been processed, thus not needed anymore
if os.path.isfile(file_path) and f in output_frames:
os.remove(file_path)
output_frames.remove(f)
if file_path.is_file() and file in output_frames:
file_path.unlink(file)
output_frames.remove(file)

View File

@ -4,4 +4,3 @@ GPUtil
psutil
requests
tqdm
future-fstrings>=1.1.0

View File

@ -1,12 +1,10 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Video2X Upscaler
Author: K4YT3X
Date Created: December 10, 2018
Last Modified: July 9, 2019
Last Modified: July 27, 2019
Dev: SAT3LL
@ -16,22 +14,28 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3),
(C) 2018-2019 K4YT3X
"""
from avalon_framework import Avalon
# local imports
from exceptions import *
from ffmpeg import Ffmpeg
from fractions import Fraction
from image_cleaner import ImageCleaner
from tqdm import tqdm
from waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter
from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
# built-in imports
from fractions import Fraction
import copy
import os
import pathlib
import re
import shutil
import tempfile
import threading
import time
import traceback
# third-party imports
from avalon_framework import Avalon
from tqdm import tqdm
class Upscaler:
@ -58,16 +62,16 @@ class Upscaler:
self.scale_ratio = None
self.model_dir = None
self.threads = 5
self.video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
self.image_format = 'png'
self.preserve_frames = False
def create_temp_directories(self):
"""create temporary directory
"""
self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory)
self.extracted_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
Avalon.debug_info(f'Extracted frames are being saved to: {self.extracted_frames}')
self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory)
self.upscaled_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
Avalon.debug_info(f'Upscaled frames are being saved to: {self.upscaled_frames}')
def cleanup_temp_directories(self):
@ -81,7 +85,8 @@ class Upscaler:
print(f'Cleaning up cache directory: {directory}')
shutil.rmtree(directory)
except (OSError, FileNotFoundError):
pass
print(f'Unable to delete: {directory}')
traceback.print_exc()
def _check_arguments(self):
# check if arguments are valid / all necessary argument
@ -106,7 +111,7 @@ class Upscaler:
# get number of extracted frames
total_frames = 0
for directory in extracted_frames_directories:
total_frames += len([f for f in os.listdir(directory) if f[-4:] == f'.{self.image_format}'])
total_frames += len([f for f in directory.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
with tqdm(total=total_frames, ascii=True, desc='Upscaling Progress') as progress_bar:
@ -117,7 +122,7 @@ class Upscaler:
while not self.progress_bar_exit_signal:
try:
total_frames_upscaled = len([f for f in os.listdir(self.upscaled_frames) if f[-4:] == f'.{self.image_format}'])
total_frames_upscaled = len([f for f in self.upscaled_frames.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
delta = total_frames_upscaled - previous_cycle_frames
previous_cycle_frames = total_frames_upscaled
@ -163,9 +168,9 @@ class Upscaler:
progress_bar.start()
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 self.upscaled_frames.iterdir() if f.is_file()]:
renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image)
shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed))
(self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
self.progress_bar_exit_signal = True
progress_bar.join()
@ -175,7 +180,7 @@ class Upscaler:
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))]
frames = [(self.extracted_frames / f) for f in self.extracted_frames.iterdir() if f.is_file]
# if we have less images than threads,
# create only the threads necessary
@ -188,13 +193,13 @@ class Upscaler:
thread_pool = []
thread_directories = []
for thread_id in range(self.threads):
thread_directory = os.path.join(self.extracted_frames, str(thread_id))
thread_directory = self.extracted_frames / str(thread_id)
thread_directories.append(thread_directory)
# delete old directories and create new directories
if os.path.isdir(thread_directory):
if thread_directory.is_dir():
shutil.rmtree(thread_directory)
os.mkdir(thread_directory)
thread_directory.mkdir(parents=True, exist_ok=True)
# append directory path into list
thread_pool.append((thread_directory, thread_id))
@ -203,7 +208,7 @@ class Upscaler:
# until there is none left in the directory
for image in frames:
# move image
shutil.move(image, thread_pool[0][0])
image.rename(thread_pool[0][0] / image.name)
# rotate list
thread_pool = thread_pool[-1:] + thread_pool[:-1]
@ -216,21 +221,21 @@ class Upscaler:
if self.scale_ratio:
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio,
False,
False,
self.image_format,
self.upscaler_exceptions))
self.upscaled_frames,
self.scale_ratio,
False,
False,
self.image_format,
self.upscaler_exceptions))
else:
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
False,
self.scale_width,
self.scale_height,
self.image_format,
self.upscaler_exceptions))
self.upscaled_frames,
False,
self.scale_width,
self.scale_height,
self.image_format,
self.upscaler_exceptions))
# if the driver being used is waifu2x_ncnn_vulkan
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
@ -273,7 +278,6 @@ class Upscaler:
if len(self.upscaler_exceptions) != 0:
raise(self.upscaler_exceptions[0])
def run(self):
"""Main controller for Video2X
@ -286,8 +290,8 @@ class Upscaler:
self._check_arguments()
# convert paths to absolute paths
self.input_video = os.path.abspath(self.input_video)
self.output_video = os.path.abspath(self.output_video)
self.input_video = self.input_video.absolute()
self.output_video = self.output_video.absolute()
# initialize objects for ffmpeg and waifu2x-caffe
fm = Ffmpeg(self.ffmpeg_settings, self.image_format)

View File

@ -1,8 +1,6 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
"""
# -*- coding: utf-8 -*-
r"""
__ __ _ _ ___ __ __
\ \ / / (_) | | |__ \ \ \ / /
@ -15,7 +13,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: July 9, 2019
Last Modified: July 27, 2019
Dev: BrianPetkovsek
Dev: SAT3LL
@ -43,13 +41,16 @@ 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
# local imports
from upscaler import Upscaler
# built-in imports
import argparse
import GPUtil
import glob
import json
import os
import psutil
import pathlib
import re
import shutil
import sys
@ -57,7 +58,12 @@ import tempfile
import time
import traceback
VERSION = '2.8.1'
# third-party imports
from avalon_framework import Avalon
import GPUtil
import psutil
VERSION = '2.9.0'
# each thread might take up to 2.5 GB during initialization.
# (system memory, not to be confused with GPU memory)
@ -76,16 +82,16 @@ def process_arguments():
# video options
file_options = parser.add_argument_group('File Options')
file_options.add_argument('-i', '--input', help='source video file/directory', action='store')
file_options.add_argument('-o', '--output', help='output video file/directory', action='store')
file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store')
file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store')
# 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='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'])
upscaler_options.add_argument('-y', '--model_dir', 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('-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=os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'video2x.json'))
upscaler_options.add_argument('-c', '--config', type=pathlib.Path, help='video2x config file location', action='store', default=pathlib.Path(sys.argv[0]).parent.absolute() / 'video2x.json')
upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true')
# scaling options
@ -103,12 +109,16 @@ def process_arguments():
def print_logo():
print('__ __ _ _ ___ __ __')
print('\\ \\ / / (_) | | |__ \\ \\ \\ / /')
print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /')
print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <')
print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\')
print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\')
"""print video2x logo"""
logo = r'''
__ __ _ _ ___ __ __
\ \ / / (_) | | |__ \ \ \ / /
\ \ / / _ __| | ___ ___ ) | \ V /
\ \/ / | | / _` | / _ \ / _ \ / / > <
\ / | | | (_| | | __/ | (_) | / /_ / . \
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
'''
print(logo)
print('\n Video2X Video Enlarger')
spaces = ((44 - len(f'Version {VERSION}')) // 2) * ' '
print(f'{Avalon.FM.BD}\n{spaces} Version {VERSION}\n{Avalon.FM.RST}')
@ -129,7 +139,7 @@ def check_memory():
# GPUtil requires nvidia-smi.exe to interact with GPU
if args.method == 'gpu' or args.method == 'cudnn':
if not (shutil.which('nvidia-smi') or
os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe')):
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 threads you\'re using')
@ -193,28 +203,28 @@ def absolutify_paths(config):
Returns:
dict -- configuration file dictionary
"""
current_directory = os.path.dirname(os.path.abspath(sys.argv[0]))
current_directory = pathlib.Path(sys.argv[0]).parent.absolute()
# check waifu2x-caffe path
if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
config['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(current_directory, config['waifu2x_caffe']['waifu2x_caffe_path'])
config['waifu2x_caffe']['waifu2x_caffe_path'] = current_directory / config['waifu2x_caffe']['waifu2x_caffe_path']
# check waifu2x-converter-cpp path
if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
config['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(current_directory, config['waifu2x_converter']['waifu2x_converter_path'])
config['waifu2x_converter']['waifu2x_converter_path'] = current_directory / config['waifu2x_converter']['waifu2x_converter_path']
# check waifu2x_ncnn_vulkan path
if not re.match('^[a-z]:', config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'], re.IGNORECASE):
config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(current_directory, config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'])
config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = current_directory / config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path']
# check ffmpeg path
if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
config['ffmpeg']['ffmpeg_path'] = os.path.join(current_directory, config['ffmpeg']['ffmpeg_path'])
config['ffmpeg']['ffmpeg_path'] = current_directory / config['ffmpeg']['ffmpeg_path']
# check video2x cache path
if config['video2x']['video2x_cache_directory']:
if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
config['video2x']['video2x_cache_directory'] = os.path.join(current_directory, config['video2x']['video2x_cache_directory'])
config['video2x']['video2x_cache_directory'] = current_directory / config['video2x']['video2x_cache_directory']
return config
@ -271,19 +281,19 @@ config = absolutify_paths(config)
# load waifu2x configuration
if args.driver == 'waifu2x_caffe':
waifu2x_settings = config['waifu2x_caffe']
if not os.path.isfile(waifu2x_settings['waifu2x_caffe_path']):
if not pathlib.Path(waifu2x_settings['waifu2x_caffe_path']).is_file():
Avalon.error('Specified waifu2x-caffe directory doesn\'t exist')
Avalon.error('Please check the configuration file settings')
raise FileNotFoundError(waifu2x_settings['waifu2x_caffe_path'])
elif args.driver == 'waifu2x_converter':
waifu2x_settings = config['waifu2x_converter']
if not os.path.isdir(waifu2x_settings['waifu2x_converter_path']):
if not pathlib.Path(waifu2x_settings['waifu2x_converter_path']).is_dir():
Avalon.error('Specified waifu2x-converter-cpp directory doesn\'t exist')
Avalon.error('Please check the configuration file settings')
raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path'])
elif args.driver == 'waifu2x_ncnn_vulkan':
waifu2x_settings = config['waifu2x_ncnn_vulkan']
if not os.path.isfile(waifu2x_settings['waifu2x_ncnn_vulkan_path']):
if not pathlib.Path(waifu2x_settings['waifu2x_ncnn_vulkan_path']).is_file():
Avalon.error('Specified waifu2x_ncnn_vulkan directory doesn\'t exist')
Avalon.error('Please check the configuration file settings')
raise FileNotFoundError(waifu2x_settings['waifu2x_ncnn_vulkan_path'])
@ -292,28 +302,38 @@ elif args.driver == 'waifu2x_ncnn_vulkan':
ffmpeg_settings = config['ffmpeg']
# load video2x settings
video2x_cache_directory = config['video2x']['video2x_cache_directory']
image_format = config['video2x']['image_format'].lower()
preserve_frames = config['video2x']['preserve_frames']
# create temp directories if they don't exist
if not video2x_cache_directory:
video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
# load cache directory
if isinstance(config['video2x']['video2x_cache_directory'], str):
video2x_cache_directory = pathlib.Path(config['video2x']['video2x_cache_directory'])
else:
video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
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):
Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist')
if Avalon.ask('Create directory?', default=True, batch=args.batch):
if os.mkdir(video2x_cache_directory) is None:
Avalon.info(f'{video2x_cache_directory} created')
else:
Avalon.error(f'Unable to create {video2x_cache_directory}')
Avalon.error('Aborting...')
exit(1)
if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
Avalon.error('Specified cache directory is a file/link')
raise FileExistsError('Specified cache directory is a file/link')
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:
video2x_cache_directory.mkdir(parents=True, exist_ok=True)
Avalon.info(f'{video2x_cache_directory} created')
# there can be a number of exceptions here
# PermissionError, FileExistsError, etc.
# therefore, we put a catch-them-all here
except Exception as e:
Avalon.error(f'Unable to create {video2x_cache_directory}')
Avalon.error('Aborting...')
raise e
else:
Avalon.error('Specified cache directory is a file/link')
Avalon.error('Unable to continue, exiting...')
exit(1)
raise FileNotFoundError('Could not create cache directory')
# start execution
@ -322,17 +342,17 @@ try:
begin_time = time.time()
# if input specified is a single file
if os.path.isfile(args.input):
if args.input.is_file():
# upscale single video file
Avalon.info(f'Upscaling single video file: {args.input}')
# check for input output format mismatch
if os.path.isdir(args.output):
if args.output.is_dir():
Avalon.error('Input and output path type mismatch')
Avalon.error('Input is single file but output is directory')
raise Exception('input output path type mismatch')
if not re.search('.*\..*$', args.output):
if not re.search('.*\..*$', str(args.output)):
Avalon.error('No suffix found in output file path')
Avalon.error('Suffix must be specified for FFmpeg')
raise Exception('No suffix specified')
@ -356,12 +376,12 @@ try:
upscaler.cleanup_temp_directories()
# if input specified is a directory
elif os.path.isdir(args.input):
elif args.input.is_dir():
# upscale videos in a directory
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))]:
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)
for input_video in [f for f in args.input.iterdir() if f.is_file()]:
output_video = args.output / input_video
upscaler = Upscaler(input_video=args.input / input_video, output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
# set optional options
upscaler.waifu2x_driver = args.driver

View File

@ -1,13 +1,11 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Name: Video2X Setup Script
Author: K4YT3X
Author: BrianPetkovsek
Date Created: November 28, 2018
Last Modified: June 26, 2019
Last Modified: July 27, 2019
Dev: SAT3LL
@ -26,14 +24,19 @@ Installation Details:
- waifu2x_ncnn_vulkan: %LOCALAPPDATA%\\video2x\\waifu2x-ncnn-vulkan
"""
# built-in imports
import argparse
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
import tempfile
import traceback
import urllib
import zipfile
# Requests doesn't come with windows, therefore
@ -43,6 +46,9 @@ import zipfile
VERSION = '1.3.0'
# global static variables
LOCALAPPDATA = pathlib.Path(os.getenv('localappdata'))
def process_arguments():
"""Processes CLI arguments
@ -104,14 +110,15 @@ class Video2xSetup:
"""
for file in self.trash:
try:
if os.path.isfile(file):
print('Deleting: {}'.format(file))
os.remove(file)
else:
print('Deleting: {}'.format(file))
if file.is_dir():
print(f'Deleting directory: {file}')
shutil.rmtree(file)
except FileNotFoundError:
pass
else:
print(f'Deleting file: {file}')
file.unlink()
except Exception:
print(f'Error deleting: {file}')
traceback.print_exc()
def _install_ffmpeg(self):
""" Install FFMPEG
@ -122,7 +129,7 @@ class Video2xSetup:
self.trash.append(ffmpeg_zip)
with zipfile.ZipFile(ffmpeg_zip) as zipf:
zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
zipf.extractall(LOCALAPPDATA / 'video2x')
def _install_waifu2x_caffe(self):
""" Install waifu2x_caffe
@ -139,7 +146,7 @@ class Video2xSetup:
self.trash.append(waifu2x_caffe_zip)
with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
zipf.extractall(LOCALAPPDATA / 'video2x')
def _install_waifu2x_converter_cpp(self):
""" Install waifu2x_caffe
@ -157,7 +164,7 @@ class Video2xSetup:
self.trash.append(waifu2x_converter_cpp_zip)
with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp'))
zipf.extractall(LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp')
def _install_waifu2x_ncnn_vulkan(self):
""" Install waifu2x-ncnn-vulkan
@ -176,12 +183,11 @@ class Video2xSetup:
# extract then move (to remove the top level directory)
with zipfile.ZipFile(waifu2x_ncnn_vulkan_zip) as zipf:
extraction_path = os.path.join(tempfile.gettempdir(), 'waifu2x-ncnn-vulkan-ext')
extraction_path = pathlib.Path(tempfile.gettempdir()) / 'waifu2x-ncnn-vulkan-ext'
zipf.extractall(extraction_path)
shutil.move(os.path.join(extraction_path, os.listdir(extraction_path)[0]), os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-ncnn-vulkan'))
shutil.move(extraction_path / os.listdir(extraction_path)[0], LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan')
self.trash.append(extraction_path)
def _generate_config(self):
""" Generate video2x config
"""
@ -190,21 +196,19 @@ class Video2xSetup:
template_dict = json.load(template)
template.close()
local_app_data = os.getenv('localappdata')
# configure only the specified drivers
if self.driver == 'all':
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'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x-ncnn-vulkan.exe')
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe'
template_dict['waifu2x_converter']['waifu2x_converter_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp'
template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe'
elif self.driver == 'waifu2x_caffe':
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe'
elif self.driver == 'waifu2x_converter':
template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
template_dict['waifu2x_converter']['waifu2x_converter_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp'
elif self.driver == 'waifu2x_ncnn_vulkan':
template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x-ncnn-vulkan.exe')
template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe'
template_dict['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x', 'ffmpeg-latest-win64-static', 'bin')
template_dict['ffmpeg']['ffmpeg_path'] = LOCALAPPDATA / 'video2x' / 'ffmpeg-latest-win64-static' / 'bin'
template_dict['video2x']['video2x_cache_directory'] = None
template_dict['video2x']['preserve_frames'] = False
@ -220,13 +224,42 @@ def download(url, save_path, chunk_size=4096):
from tqdm import tqdm
import requests
output_file = os.path.join(save_path, url.split('/')[-1])
print('Downloading: {}'.format(url))
print('Chunk size: {}'.format(chunk_size))
print('Saving to: {}'.format(output_file))
save_path = pathlib.Path(save_path)
stream = requests.get(url, stream=True)
total_size = int(stream.headers['content-length'])
# create target folder if it doesn't exist
save_path.mkdir(parents=True, exist_ok=True)
# create requests stream for steaming file
stream = requests.get(url, stream=True, allow_redirects=True)
# get file name
file_name = None
if 'content-disposition' in stream.headers:
disposition = stream.headers['content-disposition']
try:
file_name = re.findall("filename=(.+)", disposition)[0].strip('"')
except IndexError:
pass
if file_name is None:
# output_file = f'{save_path}\\{stream.url.split("/")[-1]}'
output_file = save_path / stream.url.split('/')[-1]
else:
output_file = save_path / file_name
# decode url encoding
output_file = pathlib.Path(urllib.parse.unquote(str(output_file)))
# get total size for progress bar if provided in headers
total_size = 0
if 'content-length' in stream.headers:
total_size = int(stream.headers['content-length'])
# print download information summary
print(f'Downloading: {url}')
print(f'Total size: {total_size}')
print(f'Chunk size: {chunk_size}')
print(f'Saving to: {output_file}')
# Write content into file
with open(output_file, 'wb') as output:
@ -236,6 +269,7 @@ def download(url, save_path, chunk_size=4096):
output.write(chunk)
progress_bar.update(len(chunk))
# return the full path of saved file
return output_file
@ -252,7 +286,7 @@ if __name__ == '__main__':
try:
args = process_arguments()
print('Video2X Setup Script')
print('Version: {}'.format(VERSION))
print(f'Version: {VERSION}')
# do not install pip modules if script
# is packaged in exe format

View File

@ -1,20 +1,22 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Waifu2x Caffe Driver
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: July 9, 2019
Last Modified: July 27, 2019
Description: This class is a high-level wrapper
for waifu2x-caffe.
"""
from avalon_framework import Avalon
# built-in imports
import subprocess
import threading
# third-party imports
from avalon_framework import Avalon
class Waifu2xCaffe:
"""This class communicates with waifu2x cui engine

View File

@ -1,21 +1,23 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Waifu2x Converter CPP Driver
Author: K4YT3X
Date Created: February 8, 2019
Last Modified: June 15, 2019
Last Modified: July 27, 2019
Description: This class is a high-level wrapper
for waifu2x-converter-cpp.
"""
from avalon_framework import Avalon
import os
# built-in imports
import pathlib
import subprocess
import threading
# third-party imports
from avalon_framework import Avalon
class Waifu2xConverter:
"""This class communicates with waifu2x cui engine
@ -62,8 +64,7 @@ class Waifu2xConverter:
# models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically
if self.waifu2x_settings['model-dir'] is None:
self.waifu2x_settings['model-dir'] = os.path.join(self.waifu2x_settings['waifu2x_converter_path'],
'models_rgb')
self.waifu2x_settings['model-dir'] = pathlib.Path(self.waifu2x_settings['waifu2x_converter_path']) / 'models_rgb'
# print thread start message
self.print_lock.acquire()
@ -79,7 +80,7 @@ class Waifu2xConverter:
# the key doesn't need to be passed in this case
if key == 'waifu2x_converter_path':
execute.append(os.path.join(str(value), 'waifu2x-converter-cpp.exe'))
execute.append(pathlib.Path(str(value)) / 'waifu2x-converter-cpp.exe')
# null or None means that leave this option out (keep default)
elif value is None or value is False:

View File

@ -1,23 +1,25 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
# -*- coding: utf-8 -*-
"""
Name: Waifu2x NCNN Vulkan Driver
Author: SAT3LL
Date Created: June 26, 2019
Last Modified: June 26, 2019
Last Modified: July 27, 2019
Dev: K4YT3X
Description: This class is a high-level wrapper
for waifu2x_ncnn_vulkan.
"""
from avalon_framework import Avalon
# built-in imports
import os
import subprocess
import threading
# third-party imports
from avalon_framework import Avalon
class Waifu2xNcnnVulkan:
"""This class communicates with waifu2x ncnn vulkan engine