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

fixed threading bugs found by @Aidolii

Browse Source
This commit is contained in:
k4yt3x 2019-06-13 23:49:18 -04:00
parent 8b0f82e95e
commit 8e85d6a6d3
1 changed files with 13 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
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: April 28, 2019 Last Modified: June 13, 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
@ -20,6 +20,7 @@ from image_cleaner import ImageCleaner
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
import copy
import os import os
import re import re
import shutil import shutil
@ -126,7 +127,7 @@ class Upscaler:
time.sleep(1) time.sleep(1)
def _upscale_frames(self, w2): def _upscale_frames(self):
""" Upscale video frames with waifu2x-caffe """ Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted This function upscales all the frames extracted
@ -143,9 +144,14 @@ class Upscaler:
# if this thread is not empty, then an exception has occured # if this thread is not empty, then an exception has occured
self.upscaler_exceptions = [] self.upscaler_exceptions = []
# initialize waifu2x driver
if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter':
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
# it's easier to do multi-threading with waifu2x_converter # it's easier to do multi-threading with waifu2x_converter
# the number of threads can be passed directly to waifu2x_converter # the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter': if self.waifu2x_driver == 'waifu2x_converter':
w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],)) progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
progress_bar.start() progress_bar.start()
@ -197,6 +203,10 @@ class Upscaler:
# create threads and start them # create threads and start them
for thread_info in thread_pool: for thread_info in thread_pool:
# create a separate w2 instance for each thread
w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
# 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, self.image_format, self.upscaler_exceptions)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.image_format, self.upscaler_exceptions))
@ -251,14 +261,6 @@ class Upscaler:
# initialize objects for ffmpeg and waifu2x-caffe # initialize objects for ffmpeg and waifu2x-caffe
fm = Ffmpeg(self.ffmpeg_settings, self.image_format) fm = Ffmpeg(self.ffmpeg_settings, self.image_format)
# initialize waifu2x driver
if self.waifu2x_driver == 'waifu2x_caffe':
w2 = Waifu2xCaffe(self.waifu2x_settings, self.method, self.model_dir)
elif self.waifu2x_driver == 'waifu2x_converter':
w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
else:
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
# extract frames from video # extract frames from video
fm.extract_frames(self.input_video, self.extracted_frames) fm.extract_frames(self.input_video, self.extracted_frames)
@ -292,7 +294,7 @@ class Upscaler:
# upscale images one by one using waifu2x # upscale images one by one using waifu2x
Avalon.info('Starting to upscale extracted images') Avalon.info('Starting to upscale extracted images')
self._upscale_frames(w2) self._upscale_frames()
Avalon.info('Upscaling completed') Avalon.info('Upscaling completed')
# frames to Video # frames to Video