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
### 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)
- 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)

48
bin/upscaler.py
View File

@ -4,7 +4,7 @@
Name: Video2X Upscaler
Author: K4YT3X
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),
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 image_cleaner import ImageCleaner
from exceptions import *
from ffmpeg import Ffmpeg
from fractions import Fraction
from tqdm import tqdm
from waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter
from clear_image import ClearImage
import os
import re
import shutil
@ -59,12 +59,10 @@ class Upscaler:
# create temporary folder/directories
self.video2x_cache_folder = video2x_cache_folder
self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder)
self.extracted_frames = self.extracted_frames_object.name
self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
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 = self.upscaled_frames_object.name
self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames))
self.preserve_frames = preserve_frames
@ -74,10 +72,13 @@ class Upscaler:
# avalon framework cannot be used if python is shutting down
# therefore, plain print is used
if not self.preserve_frames:
print('Cleaning up cache directory: {}'.format(self.extracted_frames))
self.extracted_frames_object.cleanup()
print('Cleaning up cache directory: {}'.format(self.upscaled_frames))
self.upscaled_frames_object.cleanup()
for directory in [self.extracted_frames, self.upscaled_frames]:
try:
print('Cleaning up cache directory: {}'.format())
shutil.rmtree(directory)
except (OSError, FileNotFoundError):
pass
def _check_arguments(self):
# check if arguments are valid / all necessary argument
@ -160,6 +161,10 @@ class Upscaler:
# create a container for all 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
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:
# create thread
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:
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]
# add threads into the pool
@ -209,10 +214,10 @@ class Upscaler:
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,))
progress_bar.start()
#Create the clearer and start it
Avalon.debug_info('Starting image clearer...')
image_clear = ClearImage(self.extracted_frames,self.upscaled_frames,len(upscaler_threads))
image_clear.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:
@ -222,12 +227,15 @@ class Upscaler:
for thread in upscaler_threads:
thread.join()
#upscaling done... kill the clearer
Avalon.debug_info('Stoping image clearer...')
image_clear.stop()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.threads_exceptions) != 0:
raise(self.threads_exceptions[0])
def run(self):
"""Main controller for Video2X
@ -300,5 +308,3 @@ class Upscaler:
# migrate audio tracks and subtitles
Avalon.info('Migrating audio tracks and subtitles to upscaled video')
fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)

16
bin/video2x.py
View File

@ -13,7 +13,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller
Author: K4YT3X
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),
available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -50,7 +50,7 @@ import tempfile
import time
import traceback
VERSION = '2.6.2'
VERSION = '2.6.3'
# each thread might take up to 2.5 GB during initialization.
# (system memory, not to be confused with GPU memory)
@ -201,8 +201,19 @@ config = read_config(args.config)
# load waifu2x configuration
if args.driver == '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':
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
ffmpeg_settings = config['ffmpeg']
@ -258,6 +269,7 @@ try:
except Exception:
Avalon.error('An exception has occurred')
traceback.print_exc()
Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
finally:
# remove Video2X Cache folder
try:

7
bin/waifu2x_caffe.py
View File

@ -4,7 +4,7 @@
Name: Waifu2x Caffe Driver
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: March 19, 2019
Last Modified: March 24, 2019
Description: This class is a high-level wrapper
for waifu2x-caffe.
@ -33,7 +33,7 @@ class Waifu2xCaffe:
self.model_dir = model_dir
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
Arguments:
@ -43,6 +43,7 @@ class Waifu2xCaffe:
height {int} -- output video height
"""
try:
# overwrite config file settings
self.waifu2x_settings['input_path'] = input_folder
self.waifu2x_settings['output_path'] = output_folder
@ -86,3 +87,5 @@ class Waifu2xCaffe:
# return command execution return code
return completed_command.returncode
except Exception as e:
threads_exceptions.append(e)

7
bin/waifu2x_converter.py
View File

@ -4,7 +4,7 @@
Name: Waifu2x Converter CPP Driver
Author: K4YT3X
Date Created: February 8, 2019
Last Modified: March 19, 2019
Last Modified: March 24, 2019
Description: This class is a high-level wrapper
for waifu2x-converter-cpp.
@ -28,7 +28,7 @@ class Waifu2xConverter:
self.waifu2x_settings['model_dir'] = model_dir
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
This method executes the upscaling of extracted frames.
@ -39,6 +39,7 @@ class Waifu2xConverter:
threads {int} -- number of threads
"""
try:
# overwrite config file settings
self.waifu2x_settings['input'] = input_folder
self.waifu2x_settings['output'] = output_folder
@ -91,3 +92,5 @@ class Waifu2xConverter:
Avalon.debug_info('Executing: {}'.format(execute))
return subprocess.run(execute, check=True).returncode
except Exception as e:
threads_exceptions.append(e)