2019-06-25 23:25:27 +00:00
|
|
|
#!/usr/bin/env python3
|
2019-07-27 17:39:40 +00:00
|
|
|
# -*- coding: utf-8 -*-
|
2019-06-25 23:25:27 +00:00
|
|
|
"""
|
|
|
|
Name: Waifu2x NCNN Vulkan Driver
|
2019-06-26 03:36:03 +00:00
|
|
|
Author: SAT3LL
|
|
|
|
Date Created: June 26, 2019
|
2019-11-15 06:09:55 +00:00
|
|
|
Last Modified: November 15, 2019
|
2019-06-26 03:36:03 +00:00
|
|
|
|
|
|
|
Dev: K4YT3X
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
Description: This class is a high-level wrapper
|
2019-06-25 23:50:29 +00:00
|
|
|
for waifu2x_ncnn_vulkan.
|
2019-06-25 23:25:27 +00:00
|
|
|
"""
|
2019-07-27 17:39:40 +00:00
|
|
|
|
|
|
|
# built-in imports
|
2019-06-26 03:36:03 +00:00
|
|
|
import os
|
2019-06-25 23:25:27 +00:00
|
|
|
import subprocess
|
|
|
|
import threading
|
2019-06-26 03:36:03 +00:00
|
|
|
|
2019-07-27 17:39:40 +00:00
|
|
|
# third-party imports
|
|
|
|
from avalon_framework import Avalon
|
|
|
|
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
class Waifu2xNcnnVulkan:
|
|
|
|
"""This class communicates with waifu2x ncnn vulkan 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.
|
|
|
|
"""
|
|
|
|
|
2019-10-20 01:52:11 +00:00
|
|
|
def __init__(self, driver_settings):
|
|
|
|
self.driver_settings = driver_settings
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
# arguments passed through command line overwrites config file values
|
|
|
|
|
2019-06-25 23:50:29 +00:00
|
|
|
# waifu2x_ncnn_vulkan can't find its own model directory if its not in the current dir
|
2019-06-25 23:25:27 +00:00
|
|
|
# so change to it
|
2019-10-20 01:52:11 +00:00
|
|
|
os.chdir(os.path.join(self.driver_settings['waifu2x_ncnn_vulkan_path'], '..'))
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
self.print_lock = threading.Lock()
|
|
|
|
|
|
|
|
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions):
|
|
|
|
"""This is the core function for WAIFU2X class
|
|
|
|
|
|
|
|
Arguments:
|
|
|
|
input_directory {string} -- source directory path
|
|
|
|
output_directory {string} -- output directory path
|
|
|
|
ratio {int} -- output video ratio
|
|
|
|
"""
|
|
|
|
|
|
|
|
try:
|
|
|
|
# overwrite config file settings
|
2019-10-20 01:52:11 +00:00
|
|
|
self.driver_settings['i'] = input_directory
|
|
|
|
self.driver_settings['o'] = output_directory
|
2019-11-15 06:09:55 +00:00
|
|
|
self.driver_settings['s'] = int(scale_ratio)
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
# print thread start message
|
|
|
|
self.print_lock.acquire()
|
|
|
|
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started')
|
|
|
|
self.print_lock.release()
|
|
|
|
|
2019-08-03 23:45:22 +00:00
|
|
|
# list to be executed
|
|
|
|
# initialize the list with waifu2x binary path as the first element
|
2019-10-20 01:52:11 +00:00
|
|
|
execute = [str(self.driver_settings['path'])]
|
2019-08-03 23:45:22 +00:00
|
|
|
|
2019-10-20 01:52:11 +00:00
|
|
|
for key in self.driver_settings.keys():
|
2019-08-03 23:45:22 +00:00
|
|
|
|
2019-10-20 01:52:11 +00:00
|
|
|
value = self.driver_settings[key]
|
2019-08-03 23:45:22 +00:00
|
|
|
|
2019-08-04 01:29:13 +00:00
|
|
|
# is executable key or null or None means that leave this option out (keep default)
|
2019-10-20 01:52:11 +00:00
|
|
|
if key == 'path' or value is None or value is False:
|
2019-07-12 22:12:08 +00:00
|
|
|
continue
|
|
|
|
else:
|
2019-08-04 01:29:13 +00:00
|
|
|
if len(key) == 1:
|
|
|
|
execute.append(f'-{key}')
|
|
|
|
else:
|
|
|
|
execute.append(f'--{key}')
|
2019-07-12 22:12:08 +00:00
|
|
|
execute.append(str(value))
|
|
|
|
|
|
|
|
Avalon.debug_info(f'Executing: {execute}')
|
2019-08-04 01:29:13 +00:00
|
|
|
completed_command = subprocess.run(execute, check=True)
|
2019-06-25 23:25:27 +00:00
|
|
|
|
|
|
|
# print thread exiting message
|
|
|
|
self.print_lock.acquire()
|
|
|
|
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting')
|
|
|
|
self.print_lock.release()
|
|
|
|
|
2019-08-04 01:29:13 +00:00
|
|
|
# return command execution return code
|
|
|
|
return completed_command.returncode
|
2019-06-25 23:25:27 +00:00
|
|
|
except Exception as e:
|
|
|
|
upscaler_exceptions.append(e)
|