def _upscale_frames(self):
""" 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
"""
# initialize waifu2x driver
if self.driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(
_('Unrecognized driver: {}').format(self.driver))
# list all images in the extracted frames
frames = [(self.extracted_frames / f)
for f in self.extracted_frames.iterdir() if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = self.extracted_frames / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.driver in [
'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan',
'srmd_ncnn_vulkan'
]:
process_directories = [self.extracted_frames]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = process_directories[
-1:] + process_directories[:-1]
# create driver processes and start them
for process_directory in process_directories:
self.process_pool.append(
self.driver_object.upscale(process_directory,
self.upscaled_frames))
# start progress bar in a different thread
Avalon.debug_info(_('Starting progress monitor'))
self.progress_monitor = ProgressMonitor(self, process_directories)
self.progress_monitor.start()
# create the clearer and start it
Avalon.debug_info(_('Starting upscaled image cleaner'))
self.image_cleaner = ImageCleaner(self.extracted_frames,
self.upscaled_frames,
len(self.process_pool))
self.image_cleaner.start()
# wait for all process to exit
try:
self._wait()
except (Exception, KeyboardInterrupt, SystemExit) as e:
# cleanup
Avalon.debug_info(_('Killing progress monitor'))
self.progress_monitor.stop()
Avalon.debug_info(_('Killing upscaled image cleaner'))
self.image_cleaner.stop()
raise e
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.driver == 'waifu2x_converter_cpp':
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}', str(image.name))
(self.upscaled_frames / image).rename(self.upscaled_frames /
renamed)
# upscaling done, kill helper threads
Avalon.debug_info(_('Killing progress monitor'))
self.progress_monitor.stop()
Avalon.debug_info(_('Killing upscaled image cleaner'))
self.image_cleaner.stop()
def _upscale_frames(self):
""" 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
"""
# progress bar process exit signal
self.progress_bar_exit_signal = False
# initialize waifu2x driver
if self.waifu2x_driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(f'Unrecognized driver: {self.waifu2x_driver}')
# create a container for all upscaler processes
upscaler_processes = []
# list all images in the extracted frames
frames = [(self.extracted_frames / f) for f in self.extracted_frames.iterdir() if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = self.extracted_frames / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.waifu2x_driver in ['waifu2x_converter', 'anime4k'] :
process_directories = [self.extracted_frames]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = process_directories[-1:] + process_directories[:-1]
# create threads and start them
for process_directory in process_directories:
# if the driver being used is waifu2x-caffe
if self.waifu2x_driver == 'waifu2x_caffe':
driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.method, self.model_dir, self.bit_depth)
if self.scale_ratio:
upscaler_processes.append(driver.upscale(process_directory,
self.upscaled_frames,
self.scale_ratio,
False,
False,
self.image_format))
else:
upscaler_processes.append(driver.upscale(process_directory,
self.upscaled_frames,
False,
self.scale_width,
self.scale_height,
self.image_format))
# if the driver being used is waifu2x-converter-cpp
elif self.waifu2x_driver == 'waifu2x_converter':
driver = Waifu2xConverter(self.driver_settings, self.model_dir)
upscaler_processes.append(driver.upscale(process_directory,
self.upscaled_frames,
self.scale_ratio,
self.processes,
self.image_format))
# if the driver being used is waifu2x-ncnn-vulkan
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings))
upscaler_processes.append(driver.upscale(process_directory,
self.upscaled_frames,
self.scale_ratio))
# if the driver being used is anime4k
elif self.waifu2x_driver == 'anime4k':
driver = Anime4k(copy.deepcopy(self.driver_settings))
upscaler_processes += driver.upscale(process_directory,
self.upscaled_frames,
self.scale_ratio,
self.processes)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar, args=(process_directories,))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_processes))
image_cleaner.start()
# wait for all process to exit
try:
Avalon.debug_info('Main process waiting for subprocesses to exit')
for process in upscaler_processes:
Avalon.debug_info(f'Subprocess {process.pid} exited with code {process.wait()}')
except (KeyboardInterrupt, SystemExit):
Avalon.warning('Exit signal received')
Avalon.warning('Killing processes')
for process in upscaler_processes:
process.terminate()
# cleanup and exit with exit code 1
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
sys.exit(1)
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.waifu2x_driver == 'waifu2x_converter':
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}', str(image.name))
(self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
# pass exit signal to progress bar thread
self.progress_bar_exit_signal = True
def _upscale_frames(self, input_directory: pathlib.Path,
output_directory: pathlib.Path):
"""Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Args:
input_directory (pathlib.Path): directory containing frames to upscale
output_directory (pathlib.Path): directory which upscaled frames should be exported to
Raises:
UnrecognizedDriverError: raised when the given driver is not recognized
e: re-raised exception after an exception has been captured and finished processing in this scope
"""
# initialize waifu2x driver
if self.driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(
_("Unrecognized driver: {}").format(self.driver))
# list all images in the extracted frames
frames = [(input_directory / f) for f in input_directory.iterdir()
if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = input_directory / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.driver in [
"waifu2x_converter_cpp",
"waifu2x_ncnn_vulkan",
"srmd_ncnn_vulkan",
"realsr_ncnn_vulkan",
"anime4kcpp",
]:
process_directories = [input_directory]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = (process_directories[-1:] +
process_directories[:-1])
# create driver processes and start them
for process_directory in process_directories:
self.process_pool.append(
self.driver_object.upscale(process_directory,
output_directory))
# start progress bar in a different thread
Avalon.debug_info(_("Starting progress monitor"))
self.progress_monitor = ProgressMonitor(self, process_directories)
self.progress_monitor.start()
# create the clearer and start it
Avalon.debug_info(_("Starting upscaled image cleaner"))
self.image_cleaner = ImageCleaner(input_directory, output_directory,
len(self.process_pool))
self.image_cleaner.start()
# wait for all process to exit
try:
self._wait()
except (Exception, KeyboardInterrupt, SystemExit) as e:
# cleanup
Avalon.debug_info(_("Killing progress monitor"))
self.progress_monitor.stop()
Avalon.debug_info(_("Killing upscaled image cleaner"))
self.image_cleaner.stop()
raise e
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.driver == "waifu2x_converter_cpp":
for image in [
f for f in output_directory.iterdir() if f.is_file()
]:
renamed = re.sub(
f"_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}",
f".{self.extracted_frame_format}",
str(image.name),
)
(output_directory / image).rename(output_directory / renamed)
# upscaling done, kill helper threads
Avalon.debug_info(_("Killing progress monitor"))
self.progress_monitor.stop()
Avalon.debug_info(_("Killing upscaled image cleaner"))
self.image_cleaner.stop()
def _upscale_frames(self):
""" 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
"""
# progress bar thread exit signal
self.progress_bar_exit_signal = False
# create a container for exceptions in threads
# if this thread is not empty, then an exception has occured
self.upscaler_exceptions = []
# initialize waifu2x driver
drivers = AVAILABLE_DRIVERS
if self.waifu2x_driver not in drivers:
raise UnrecognizedDriverError(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
# it's easier to do multi-threading with waifu2x_converter
# the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter':
w2 = Waifu2xConverter(self.driver_settings, self.model_dir)
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
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 self.upscaled_frames.iterdir() if f.is_file()]:
renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', str(image))
(self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
self.progress_bar_exit_signal = True
progress_bar.join()
return
# drivers that are to be multi-threaded by video2x
else:
# create a container for all upscaler threads
upscaler_threads = []
# list all images in the extracted frames
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
if len(frames) < self.threads:
self.threads = len(frames)
# create a directory for each thread and append directory
# name into a list
thread_pool = []
thread_directories = []
for thread_id in range(self.threads):
thread_directory = self.extracted_frames / str(thread_id)
thread_directories.append(thread_directory)
# delete old directories and create new directories
if thread_directory.is_dir():
shutil.rmtree(thread_directory)
thread_directory.mkdir(parents=True, exist_ok=True)
# append directory path into list
thread_pool.append((thread_directory, thread_id))
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(thread_pool[0][0] / image.name)
# rotate list
thread_pool = thread_pool[-1:] + thread_pool[:-1]
# create threads and start them
for thread_info in thread_pool:
# create a separate w2 instance for each thread
if self.waifu2x_driver == 'waifu2x_caffe':
w2 = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.method, self.model_dir, self.bit_depth)
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))
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))
# if the driver being used is waifu2x_ncnn_vulkan
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings))
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio,
self.upscaler_exceptions))
# if the driver being used is anime4k
elif self.waifu2x_driver == 'anime4k':
w2 = Anime4k(copy.deepcopy(self.driver_settings))
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio,
self.upscaler_exceptions))
# create thread
thread.name = thread_info[1]
# add threads into the pool
upscaler_threads.append(thread)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
image_cleaner.start()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.upscaler_exceptions) != 0:
raise(self.upscaler_exceptions[0])
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
"""
# progress bar thread exit signal
self.progress_bar_exit_signal = False
# create a container for exceptions in threads
# if this thread is not empty, then an exception has occured
self.upscaler_exceptions = []
# it's easier to do multi-threading with waifu2x_converter
# the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter':
progress_bar = threading.Thread(target=self._progress_bar,
args=([self.extracted_frames], ))
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))
]:
renamed = re.sub(
'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{}'.format(
self.image_format), '.{}'.format(self.image_format),
image)
shutil.move('{}\\{}'.format(self.upscaled_frames, image),
'{}\\{}'.format(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True
progress_bar.join()
return
# 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 = []
thread_folders = []
for thread_id in range(self.threads):
thread_folder = '{}\\{}'.format(self.extracted_frames,
str(thread_id))
thread_folders.append(thread_folder)
# 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
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))
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))
thread.name = thread_info[1]
# add threads into the pool
upscaler_threads.append(thread)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar,
args=(thread_folders, ))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames,
self.upscaled_frames,
len(upscaler_threads))
image_cleaner.start()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.upscaler_exceptions) != 0:
raise (self.upscaler_exceptions[0])