Merge pull request #55 from BrianPetkovsek/master

Fixed execution problem & added a file clearer
This commit is contained in:
K4YT3X 2019-03-24 22:03:09 -04:00 committed by GitHub
commit ecda3d6990
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 69 additions and 7 deletions

54
bin/clear_image.py Normal file
View File

@ -0,0 +1,54 @@
#!/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)

View File

@ -19,6 +19,7 @@ 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
@ -208,6 +209,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
Avalon.debug_info('Starting image clearer...')
image_clear = ClearImage(self.extracted_frames,self.upscaled_frames,len(upscaler_threads))
image_clear.start()
# start all threads # start all threads
for thread in upscaler_threads: for thread in upscaler_threads:
thread.start() thread.start()
@ -215,6 +221,10 @@ 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
Avalon.debug_info('Stoping image clearer...')
image_clear.stop()
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
@ -290,3 +300,5 @@ 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)

View File

@ -61,18 +61,14 @@ class Waifu2xCaffe:
# list to be executed # list to be executed
execute = [] execute = []
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]
# the key doesn't need to be passed in this case #is executable key or null or None means that leave this option out (keep default)
if key == 'waifu2x_caffe_path': if key == 'waifu2x_caffe_path' or value is None or value is False:
execute.append(str(value))
# null or None means that leave this option out (keep default)
elif value is None or value is False:
continue continue
else: else:
if len(key) == 1: if len(key) == 1:
execute.append('-{}'.format(key)) execute.append('-{}'.format(key))