OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-28 23:19:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #1 from K4YT3X/master

Browse Source 
Pull latest updates
This commit is contained in:
Brian Petkovsek 2019-03-25 02:14:50 -04:00 committed by GitHub
commit b984977227
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 210 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@ -16,6 +16,12 @@ Component names that are *italicized* can be automatically downloaded and config
## Recent Changes ## Recent Changes
### 2.6.3 (March 24, 2019)
- Added image cleaner by @BrianPetkovsek which removes upscaled frames.
- Fixed some PEP8 issues.
- Exceptions in waifu2x are now caught, and script will now stop on waifu2x error instead of keep going on to FFMPEG.
### 2.6.2 (March 19, 2019) ### 2.6.2 (March 19, 2019)
- Removed `--model_dir` verification due to the rapidly evolving number of models added. - Removed `--model_dir` verification due to the rapidly evolving number of models added.

54
bin/clear_image.py
View File

@ -1,54 +0,0 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Name: Waifu2x Image clearer
Author: BrianPetkovsek
Date Created: March 24, 2019
Last Modified: March 25, 2019
Description: This class is to remove the
downscaled image files when upscale is finished
from waifu2x-caffe.
"""
from threading import Thread
from time import sleep
import os
class ClearImage(Thread):
def __init__(self, input_folder, output_folder,num_threads):
Thread.__init__(self)
self.input_folder = input_folder
self.output_folder = output_folder
self.num_threads = num_threads
self.running = False
def run(self):
self.running = True
while(self.running):
self.removeFrames()
#delay in 1 second intrvals for stop trigger
i=0
while self.running and i<20:
i+=1
sleep(1)
def stop(self):
self.running = False
self.join()
def removeFrames(self):
# list all images in the extracted frames
output_frames = [f for f in os.listdir(self.output_folder) if os.path.isfile(os.path.join(self.output_folder, f))]
# compare and remove frames downscaled images that finished being upscaled
for i in range(self.num_threads):
dir_path = os.path.join(self.input_folder,str(i))
for f in os.listdir(dir_path):
file_path = os.path.join(dir_path, f)
if os.path.isfile(file_path) and f in output_frames:
os.remove(file_path)
output_frames.remove(f)

76
bin/image_cleaner.py Normal file
View File

@ -0,0 +1,76 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Name: Video2X Image Cleaner
Author: BrianPetkovsek
Author: K4YT3X
Date Created: March 24, 2019
Last Modified: March 24, 2019
Description: This class is to remove the extracted frames
that have already been upscaled.
"""
import os
import threading
import time
class ImageCleaner(threading.Thread):
""" Video2X Image Cleaner
This class creates an object that keeps track of extracted
frames that has already been upscaled and are not needed
anymore. It then deletes them to save disk space.
Extends:
threading.Thread
"""
def __init__(self, input_folder, output_folder, num_threads):
threading.Thread.__init__(self)
self.input_folder = input_folder
self.output_folder = output_folder
self.num_threads = num_threads
self.running = False
def run(self):
""" Run image cleaner
"""
self.running = True
while self.running:
self.remove_upscaled_frames()
time.sleep(1)
def stop(self):
""" Stop the image cleaner
"""
self.running = False
self.join()
def remove_upscaled_frames(self):
""" remove frames that have already been upscaled
This method compares the files in the extracted frames
folder with the upscaled frames folder, and removes
the frames that has already been upscaled.
"""
# list all images in the extracted frames
output_frames = [f for f in os.listdir(self.output_folder) if os.path.isfile(os.path.join(self.output_folder, f))]
# compare and remove frames downscaled images that finished being upscaled
# within each thread's extracted frames folder
for i in range(self.num_threads):
dir_path = os.path.join(self.input_folder, str(i))
# for each file within all the folders
for f in os.listdir(dir_path):
file_path = os.path.join(dir_path, f)
# if file also exists in the output folder, 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)

52
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: March 19, 2019 Last Modified: March 24, 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
@ -13,13 +13,13 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3),
""" """
from avalon_framework import Avalon from avalon_framework import Avalon
from image_cleaner import ImageCleaner
from exceptions import * from exceptions import *
from ffmpeg import Ffmpeg from ffmpeg import Ffmpeg
from fractions import Fraction from fractions import Fraction
from tqdm import tqdm from tqdm import tqdm
from waifu2x_caffe import Waifu2xCaffe from waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter from waifu2x_converter import Waifu2xConverter
from clear_image import ClearImage
import os import os
import re import re
import shutil import shutil
@ -59,12 +59,10 @@ class Upscaler:
# create temporary folder/directories # create temporary folder/directories
self.video2x_cache_folder = video2x_cache_folder self.video2x_cache_folder = video2x_cache_folder
self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
self.extracted_frames = self.extracted_frames_object.name
Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames)) Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames))
self.upscaled_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
self.upscaled_frames = self.upscaled_frames_object.name
Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames)) Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames))
self.preserve_frames = preserve_frames self.preserve_frames = preserve_frames
@ -74,10 +72,13 @@ class Upscaler:
# avalon framework cannot be used if python is shutting down # avalon framework cannot be used if python is shutting down
# therefore, plain print is used # therefore, plain print is used
if not self.preserve_frames: if not self.preserve_frames:
print('Cleaning up cache directory: {}'.format(self.extracted_frames))
self.extracted_frames_object.cleanup() for directory in [self.extracted_frames, self.upscaled_frames]:
print('Cleaning up cache directory: {}'.format(self.upscaled_frames)) try:
self.upscaled_frames_object.cleanup() print('Cleaning up cache directory: {}'.format())
shutil.rmtree(directory)
except (OSError, FileNotFoundError):
pass
def _check_arguments(self): def _check_arguments(self):
# check if arguments are valid / all necessary argument # check if arguments are valid / all necessary argument
@ -160,6 +161,10 @@ class Upscaler:
# create a container for all upscaler threads # create a container for all upscaler threads
upscaler_threads = [] upscaler_threads = []
# create a container for exceptions in threads
# if this thread is not empty, then an exception has occured
self.threads_exceptions = []
# list all images in the extracted frames # 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 = [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))]
@ -197,9 +202,9 @@ class Upscaler:
for thread_info in thread_pool: for thread_info in thread_pool:
# create thread # create thread
if self.scale_ratio: if self.scale_ratio:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.threads_exceptions))
else: else:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.threads_exceptions))
thread.name = thread_info[1] thread.name = thread_info[1]
# add threads into the pool # add threads into the pool
@ -209,11 +214,11 @@ class Upscaler:
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,)) progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,))
progress_bar.start() progress_bar.start()
#Create the clearer and start it # create the clearer and start it
Avalon.debug_info('Starting image clearer...') Avalon.debug_info('Starting upscaled image cleaner')
image_clear = ClearImage(self.extracted_frames,self.upscaled_frames,len(upscaler_threads)) image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
image_clear.start() image_cleaner.start()
# start all threads # start all threads
for thread in upscaler_threads: for thread in upscaler_threads:
thread.start() thread.start()
@ -221,13 +226,16 @@ class Upscaler:
# wait for threads to finish # wait for threads to finish
for thread in upscaler_threads: for thread in upscaler_threads:
thread.join() thread.join()
#upscaling done... kill the clearer # upscaling done, kill the clearer
Avalon.debug_info('Stoping image clearer...') Avalon.debug_info('Killing upscaled image cleaner')
image_clear.stop() image_cleaner.stop()
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
if len(self.threads_exceptions) != 0:
raise(self.threads_exceptions[0])
def run(self): def run(self):
"""Main controller for Video2X """Main controller for Video2X
@ -300,5 +308,3 @@ class Upscaler:
# migrate audio tracks and subtitles # migrate audio tracks and subtitles
Avalon.info('Migrating audio tracks and subtitles to upscaled video') Avalon.info('Migrating audio tracks and subtitles to upscaled video')
fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames) fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)

18
bin/video2x.py
View File

@ -13,7 +13,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller Name: Video2X Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: March 19, 2019 Last Modified: March 24, 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
@ -50,7 +50,7 @@ import tempfile
import time import time
import traceback import traceback
VERSION = '2.6.2' VERSION = '2.6.3'
# each thread might take up to 2.5 GB during initialization. # each thread might take up to 2.5 GB during initialization.
# (system memory, not to be confused with GPU memory) # (system memory, not to be confused with GPU memory)
@ -120,7 +120,7 @@ def check_memory():
if not os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'): if not 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')
else: else:
try: try:
# "0" is GPU ID. Both waifu2x drivers use the first GPU available, therefore only 0 makes sense # "0" is GPU ID. Both waifu2x drivers use the first GPU available, therefore only 0 makes sense
@ -201,8 +201,19 @@ config = read_config(args.config)
# load waifu2x configuration # load waifu2x configuration
if args.driver == 'waifu2x_caffe': if args.driver == 'waifu2x_caffe':
waifu2x_settings = config['waifu2x_caffe'] waifu2x_settings = config['waifu2x_caffe']
if not os.path.isfile(waifu2x_settings['waifu2x_caffe_path']):
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': elif args.driver == 'waifu2x_converter':
waifu2x_settings = config['waifu2x_converter'] waifu2x_settings = config['waifu2x_converter']
if not os.path.isdir(waifu2x_settings['waifu2x_converter_path']):
Avalon.error('Specified waifu2x-conver-cpp directory doesn\'t exist')
Avalon.error('Please check the configuration file settings')
raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path'])
# check if waifu2x path is valid
# read FFMPEG configuration # read FFMPEG configuration
ffmpeg_settings = config['ffmpeg'] ffmpeg_settings = config['ffmpeg']
@ -258,6 +269,7 @@ try:
except Exception: except Exception:
Avalon.error('An exception has occurred') Avalon.error('An exception has occurred')
traceback.print_exc() traceback.print_exc()
Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
finally: finally:
# remove Video2X Cache folder # remove Video2X Cache folder
try: try:

75
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: March 19, 2019 Last Modified: March 24, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-caffe. for waifu2x-caffe.
@ -33,7 +33,7 @@ class Waifu2xCaffe:
self.model_dir = model_dir self.model_dir = model_dir
self.print_lock = threading.Lock() self.print_lock = threading.Lock()
def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height): def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, threads_exceptions):
"""This is the core function for WAIFU2X class """This is the core function for WAIFU2X class
Arguments: Arguments:
@ -43,46 +43,49 @@ class Waifu2xCaffe:
height {int} -- output video height height {int} -- output video height
""" """
# overwrite config file settings try:
self.waifu2x_settings['input_path'] = input_folder # overwrite config file settings
self.waifu2x_settings['output_path'] = output_folder self.waifu2x_settings['input_path'] = input_folder
self.waifu2x_settings['output_path'] = output_folder
if scale_ratio: if scale_ratio:
self.waifu2x_settings['scale_ratio'] = scale_ratio self.waifu2x_settings['scale_ratio'] = scale_ratio
elif scale_width and scale_height: elif scale_width and scale_height:
self.waifu2x_settings['scale_width'] = scale_width self.waifu2x_settings['scale_width'] = scale_width
self.waifu2x_settings['scale_height'] = scale_height self.waifu2x_settings['scale_height'] = scale_height
# print thread start message # print thread start message
self.print_lock.acquire() self.print_lock.acquire()
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))
self.print_lock.release() self.print_lock.release()
# list to be executed # list to be executed
execute = [] execute = []
execute.append(self.waifu2x_settings['waifu2x_caffe_path']) execute.append(self.waifu2x_settings['waifu2x_caffe_path'])
for key in self.waifu2x_settings.keys(): for key in self.waifu2x_settings.keys():
value = self.waifu2x_settings[key] value = self.waifu2x_settings[key]
#is executable key or null or None means that leave this option out (keep default) # is executable key or null or None means that leave this option out (keep default)
if key == 'waifu2x_caffe_path' or value is None or value is False: if key == 'waifu2x_caffe_path' or value is None or value is False:
continue continue
else:
if len(key) == 1:
execute.append('-{}'.format(key))
else: else:
execute.append('--{}'.format(key)) if len(key) == 1:
execute.append(str(value)) execute.append('-{}'.format(key))
else:
Avalon.debug_info('Executing: {}'.format(execute)) execute.append('--{}'.format(key))
completed_command = subprocess.run(execute, check=True) execute.append(str(value))
# print thread exiting message Avalon.debug_info('Executing: {}'.format(execute))
self.print_lock.acquire() completed_command = subprocess.run(execute, check=True)
Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name))
self.print_lock.release()
# return command execution return code # print thread exiting message
return completed_command.returncode self.print_lock.acquire()
Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name))
self.print_lock.release()
# return command execution return code
return completed_command.returncode
except Exception as e:
threads_exceptions.append(e)

87
bin/waifu2x_converter.py
View File

@ -4,7 +4,7 @@
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: March 19, 2019 Last Modified: March 24, 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.
@ -28,7 +28,7 @@ class Waifu2xConverter:
self.waifu2x_settings['model_dir'] = model_dir self.waifu2x_settings['model_dir'] = model_dir
self.print_lock = threading.Lock() self.print_lock = threading.Lock()
def upscale(self, input_folder, output_folder, scale_ratio, jobs): def upscale(self, input_folder, output_folder, scale_ratio, jobs, threads_exceptions):
""" Waifu2x Converter Driver Upscaler """ Waifu2x Converter Driver Upscaler
This method executes the upscaling of extracted frames. This method executes the upscaling of extracted frames.
@ -39,55 +39,58 @@ class Waifu2xConverter:
threads {int} -- number of threads threads {int} -- number of threads
""" """
# overwrite config file settings try:
self.waifu2x_settings['input'] = input_folder # overwrite config file settings
self.waifu2x_settings['output'] = output_folder self.waifu2x_settings['input'] = input_folder
self.waifu2x_settings['output'] = output_folder
# temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109
self.waifu2x_settings['i'] = input_folder self.waifu2x_settings['i'] = input_folder
self.waifu2x_settings['o'] = output_folder self.waifu2x_settings['o'] = output_folder
self.waifu2x_settings['input'] = None self.waifu2x_settings['input'] = None
self.waifu2x_settings['output'] = None self.waifu2x_settings['output'] = None
self.waifu2x_settings['scale_ratio'] = scale_ratio self.waifu2x_settings['scale_ratio'] = scale_ratio
self.waifu2x_settings['jobs'] = jobs self.waifu2x_settings['jobs'] = jobs
# 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'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path']) self.waifu2x_settings['model_dir'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path'])
# print thread start message # print thread start message
self.print_lock.acquire() self.print_lock.acquire()
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))
self.print_lock.release() self.print_lock.release()
# list to be executed # list to be executed
execute = [] execute = []
for key in self.waifu2x_settings.keys(): for key in self.waifu2x_settings.keys():
value = self.waifu2x_settings[key] value = self.waifu2x_settings[key]
# 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('{}\\waifu2x-converter-cpp.exe'.format(str(value))) execute.append('{}\\waifu2x-converter-cpp.exe'.format(str(value)))
# 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:
continue
else:
if len(key) == 1:
execute.append('-{}'.format(key))
else:
execute.append('--{}'.format(key))
# true means key is an option
if value is True:
continue continue
else:
if len(key) == 1:
execute.append('-{}'.format(key))
else:
execute.append('--{}'.format(key))
execute.append(str(value)) # true means key is an option
if value is True:
Avalon.debug_info('Executing: {}'.format(execute)) continue
return subprocess.run(execute, check=True).returncode
execute.append(str(value))
Avalon.debug_info('Executing: {}'.format(execute))
return subprocess.run(execute, check=True).returncode
except Exception as e:
threads_exceptions.append(e)