mirror of
https://github.com/k4yt3x/video2x.git
synced 2024-12-29 16:09:10 +00:00
Merge remote-tracking branch 'origin/master'
This commit is contained in:
k4yt3x
commit
685d894bef
@ -108,12 +108,13 @@ class Upscaler:
|
||||
and the output directory. This is originally
|
||||
suggested by @ArmandBernard.
|
||||
"""
|
||||
# get number of extracted frames
|
||||
total_frames = 0
|
||||
for directory in extracted_frames_directories:
|
||||
total_frames += len([f for f in directory.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
|
||||
|
||||
with tqdm(total=total_frames, ascii=True, desc='Upscaling Progress') as progress_bar:
|
||||
# get number of extracted frames
|
||||
self.total_frames = 0
|
||||
for directory in extracted_frames_directories:
|
||||
self.total_frames += len([f for f in directory.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
|
||||
|
||||
with tqdm(total=self.total_frames, ascii=True, desc='Upscaling Progress') as progress_bar:
|
||||
|
||||
# tqdm update method adds the value to the progress
|
||||
# bar instead of setting the value. Therefore, a delta
|
||||
@ -122,12 +123,12 @@ class Upscaler:
|
||||
while not self.progress_bar_exit_signal:
|
||||
|
||||
try:
|
||||
total_frames_upscaled = len([f for f in self.upscaled_frames.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
|
||||
delta = total_frames_upscaled - previous_cycle_frames
|
||||
previous_cycle_frames = total_frames_upscaled
|
||||
self.total_frames_upscaled = len([f for f in self.upscaled_frames.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
|
||||
delta = self.total_frames_upscaled - previous_cycle_frames
|
||||
previous_cycle_frames = self.total_frames_upscaled
|
||||
|
||||
# if upscaling is finished
|
||||
if total_frames_upscaled >= total_frames:
|
||||
if self.total_frames_upscaled >= self.total_frames:
|
||||
return
|
||||
|
||||
# adds the detla into the progress bar
|
||||
|
||||
374
bin/video2x_gui.py
Normal file
374
bin/video2x_gui.py
Normal file
@ -0,0 +1,374 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Name: Video2x GUI
|
||||
Author: K4YT3X
|
||||
Date Created: July 27, 2019
|
||||
Last Modified: July 27, 2019
|
||||
|
||||
Description: GUI for Video2X
|
||||
"""
|
||||
|
||||
# local imports
|
||||
from exceptions import *
|
||||
from upscaler import Upscaler
|
||||
|
||||
# built-in imports
|
||||
from tkinter import *
|
||||
from tkinter import messagebox
|
||||
from tkinter import ttk
|
||||
from tkinter.filedialog import *
|
||||
import json
|
||||
import pathlib
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import tkinter as tk
|
||||
|
||||
VERSION = '1.0.0'
|
||||
|
||||
|
||||
# global static variables
|
||||
|
||||
AVAILABLE_METHODS = {'GPU': 'gpu',
|
||||
'CUDNN': 'cudnn',
|
||||
'CPU': 'cpu'
|
||||
}
|
||||
|
||||
AVAILABLE_DRIVERS = {
|
||||
'Waifu2X Caffe': 'waifu2x_caffe',
|
||||
'Waifu2X Converter CPP': 'waifu2x_converter',
|
||||
'Waifu2x NCNN Vulkan': 'waifu2x_ncnn_vulkan'
|
||||
}
|
||||
|
||||
IMAGE_FORMATS = {'PNG', 'JPG'}
|
||||
|
||||
|
||||
class Video2xGui():
|
||||
|
||||
def __init__(self):
|
||||
|
||||
self.running = False
|
||||
|
||||
# create main window
|
||||
self.main_window = Tk()
|
||||
self.main_window.title('Video2X GUI')
|
||||
self.main_frame = Frame()
|
||||
self.main_frame.pack(fill=BOTH, expand=True)
|
||||
|
||||
# file frame
|
||||
self.file_frame = Frame(self.main_frame)
|
||||
self.file_frame.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
|
||||
# input file
|
||||
self.input_file = StringVar()
|
||||
label_text = StringVar()
|
||||
label_text.set('Input File')
|
||||
Label(self.file_frame, textvariable=label_text, relief=RIDGE, width=10).grid(row=0, column=0, padx=5, pady=5, sticky=W)
|
||||
Entry(self.file_frame, textvariable=self.input_file, width=60).grid(row=0, column=1, padx=5, pady=5, sticky=W)
|
||||
Button(self.file_frame, text='Select', command=self._select_input).grid(row=0, column=2, padx=5, pady=5, sticky=W)
|
||||
|
||||
# output file
|
||||
self.output_file = StringVar()
|
||||
label_text = StringVar()
|
||||
label_text.set('Output File')
|
||||
Label(self.file_frame, textvariable=label_text, relief=RIDGE, width=10).grid(row=1, column=0, padx=5, pady=5, sticky=W)
|
||||
Entry(self.file_frame, textvariable=self.output_file, width=60).grid(row=1, column=1, padx=5, pady=5, sticky=W)
|
||||
Button(self.file_frame, text='Select', command=self._select_output).grid(row=1, column=2, padx=5, pady=5, sticky=W)
|
||||
|
||||
# options
|
||||
self.options_frame = Frame()
|
||||
# self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
self.options_frame.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
|
||||
self.options_left = Frame(self.options_frame)
|
||||
# self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
self.options_left.grid(row=0, column=0, padx=5, pady=5, sticky=N)
|
||||
|
||||
# width
|
||||
self.width = IntVar()
|
||||
# self.width.set(1920)
|
||||
Label(self.options_left, text='Width', relief=RIDGE, width=15).grid(row=0, column=0, padx=2, pady=3)
|
||||
width_field = Entry(self.options_left, textvariable=self.width)
|
||||
width_field.grid(row=0, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# height
|
||||
self.height = IntVar()
|
||||
# self.height.set(1080)
|
||||
Label(self.options_left, text='Height', relief=RIDGE, width=15).grid(row=1, column=0, padx=2, pady=3)
|
||||
height_field = Entry(self.options_left, textvariable=self.height)
|
||||
height_field.grid(row=1, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# scale ratio
|
||||
self.scale_ratio = DoubleVar()
|
||||
# self.scale_ratio.set(2.0)
|
||||
Label(self.options_left, text='Scale Ratio', relief=RIDGE, width=15).grid(row=2, column=0, padx=2, pady=3)
|
||||
scale_ratio_field = Entry(self.options_left, textvariable=self.scale_ratio)
|
||||
scale_ratio_field.grid(row=2, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# image format
|
||||
self.image_format = StringVar(self.options_left)
|
||||
self.image_format.set('PNG')
|
||||
Label(self.options_left, text='Image Format', relief=RIDGE, width=15).grid(row=3, column=0, padx=2, pady=3)
|
||||
image_format_menu = OptionMenu(self.options_left, self.image_format, *IMAGE_FORMATS)
|
||||
image_format_menu.grid(row=3, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# options
|
||||
self.options_right = Frame(self.options_frame)
|
||||
# self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
self.options_right.grid(row=0, column=1, padx=5, pady=5, sticky=N)
|
||||
|
||||
# threads
|
||||
self.threads = IntVar()
|
||||
self.threads.set(1)
|
||||
Label(self.options_right, text='Threads', relief=RIDGE, width=15).grid(row=0, column=0, padx=2, pady=3)
|
||||
threads_field = Entry(self.options_right, textvariable=self.threads)
|
||||
threads_field.grid(row=0, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# method
|
||||
self.method = StringVar(self.options_left)
|
||||
self.method.set('GPU')
|
||||
Label(self.options_right, text='Method', relief=RIDGE, width=15).grid(row=1, column=0, padx=2, pady=3)
|
||||
method_menu = OptionMenu(self.options_right, self.method, *AVAILABLE_METHODS)
|
||||
method_menu.grid(row=1, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# driver
|
||||
self.driver = StringVar(self.options_left)
|
||||
self.driver.set('Waifu2X Caffe')
|
||||
Label(self.options_right, text='Driver', relief=RIDGE, width=15).grid(row=2, column=0, padx=2, pady=3)
|
||||
driver_menu = OptionMenu(self.options_right, self.driver, *AVAILABLE_DRIVERS)
|
||||
driver_menu.grid(row=2, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# preserve frames
|
||||
self.preserve_frames = BooleanVar(self.options_left)
|
||||
self.preserve_frames.set(True)
|
||||
Label(self.options_right, text='Preserve Frames', relief=RIDGE, width=15).grid(row=3, column=0, padx=2, pady=3)
|
||||
preserve_frames_menu = OptionMenu(self.options_right, self.preserve_frames, *{True, False})
|
||||
preserve_frames_menu.grid(row=3, column=1, padx=2, pady=3, sticky=W)
|
||||
|
||||
# progress bar
|
||||
self.progress_bar_frame = Frame()
|
||||
self.progress_bar_frame.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
|
||||
self.progress_bar = ttk.Progressbar(self.progress_bar_frame, orient='horizontal', length=100, mode='determinate')
|
||||
self.progress_bar.pack(fill=X)
|
||||
|
||||
# start button frame
|
||||
self.start_frame = Frame()
|
||||
self.start_frame.pack(fill=X, padx=5, pady=5, expand=True)
|
||||
|
||||
# start button
|
||||
self.start_button_text = StringVar()
|
||||
self.start_button_text.set('Start')
|
||||
Button(self.start_frame, textvariable=self.start_button_text, command=self._launch_upscaling, width=20).pack(side=RIGHT)
|
||||
|
||||
self.main_frame.mainloop()
|
||||
|
||||
def _launch_upscaling(self):
|
||||
|
||||
# prevent launching multiple instances
|
||||
if self.running:
|
||||
messagebox.showerror('Error', 'Video2X is already running')
|
||||
return
|
||||
|
||||
# arguments sanity check
|
||||
if self.input_file.get() == '':
|
||||
messagebox.showerror('Error', 'You must specify input video file/directory path')
|
||||
return
|
||||
if self.output_file.get() == '':
|
||||
messagebox.showerror('Error', 'You must specify output video file/directory path')
|
||||
return
|
||||
if (self.driver.get() == 'waifu2x_converter' or self.driver.get() == 'waifu2x_ncnn_vulkan') and self.width.get() and self.height.get():
|
||||
messagebox.showerror('Error', 'Waifu2x Converter CPP/NCNN accepts only scaling ratio')
|
||||
return
|
||||
if self.driver.get() == 'waifu2x_ncnn_vulkan' and (self.scale_ratio.get() > 2 or not self.scale_ratio.get().is_integer()):
|
||||
messagebox.showerror('Error', 'Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
|
||||
return
|
||||
if (self.width.get() or self.height.get()) and self.scale_ratio.get():
|
||||
messagebox.showerror('Error', 'You can only specify either scaling ratio or output width and height')
|
||||
return
|
||||
if (self.width.get() and not self.height.get()) or (not self.width.get() and self.height.get()):
|
||||
messagebox.showerror('Error', 'You must specify both width and height')
|
||||
return
|
||||
if (not self.width.get() or not self.height.get()) and not self.scale_ratio.get():
|
||||
messagebox.showerror('Error', 'You must specify either output dimensions or scaling ratio')
|
||||
return
|
||||
|
||||
upscale = threading.Thread(target=self._upscale)
|
||||
upscale.start()
|
||||
self.running = True
|
||||
self.start_button_text.set('Running')
|
||||
|
||||
def _upscale(self):
|
||||
|
||||
# read configuration file
|
||||
config = read_config('video2x.json')
|
||||
config = absolutify_paths(config)
|
||||
|
||||
input_file = pathlib.Path(self.input_file.get())
|
||||
output_file = pathlib.Path(self.output_file.get())
|
||||
driver = AVAILABLE_DRIVERS[self.driver.get()]
|
||||
|
||||
if driver == 'waifu2x_caffe':
|
||||
waifu2x_settings = config['waifu2x_caffe']
|
||||
if not pathlib.Path(waifu2x_settings['waifu2x_caffe_path']).is_file():
|
||||
messagebox.showerror('Error', 'Specified waifu2x-caffe directory doesn\'t exist\nPlease check the configuration file settings')
|
||||
raise FileNotFoundError(waifu2x_settings['waifu2x_caffe_path'])
|
||||
elif driver == 'waifu2x_converter':
|
||||
waifu2x_settings = config['waifu2x_converter']
|
||||
if not pathlib.Path(waifu2x_settings['waifu2x_converter_path']).is_dir():
|
||||
messagebox.showerror('Error', 'Specified waifu2x-converter-cpp directory doesn\'t exist\nPlease check the configuration file settings')
|
||||
elif driver == 'waifu2x_ncnn_vulkan':
|
||||
waifu2x_settings = config['waifu2x_ncnn_vulkan']
|
||||
if not pathlib.Path(waifu2x_settings['waifu2x_ncnn_vulkan_path']).is_file():
|
||||
messagebox.showerror('Error', 'Specified waifu2x_ncnn_vulkan directory doesn\'t exist\nPlease check the configuration file settings')
|
||||
raise FileNotFoundError(waifu2x_settings['waifu2x_ncnn_vulkan_path'])
|
||||
|
||||
# read FFmpeg configuration
|
||||
ffmpeg_settings = config['ffmpeg']
|
||||
# load video2x settings
|
||||
image_format = config['video2x']['image_format'].lower()
|
||||
preserve_frames = config['video2x']['preserve_frames']
|
||||
|
||||
# load cache directory
|
||||
if isinstance(config['video2x']['video2x_cache_directory'], str):
|
||||
video2x_cache_directory = pathlib.Path(config['video2x']['video2x_cache_directory'])
|
||||
else:
|
||||
video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
|
||||
|
||||
if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
|
||||
messagebox.showerror('Error', 'Specified cache directory is a file/link')
|
||||
raise FileExistsError('Specified cache directory is a file/link')
|
||||
|
||||
elif not video2x_cache_directory.exists():
|
||||
# try creating the cache directory
|
||||
if messagebox.askyesno('Question', f'Specified cache directory {video2x_cache_directory} does not exist\nCreate directory?'):
|
||||
try:
|
||||
video2x_cache_directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# there can be a number of exceptions here
|
||||
# PermissionError, FileExistsError, etc.
|
||||
# therefore, we put a catch-them-all here
|
||||
except Exception as e:
|
||||
messagebox.showerror('Error', f'Unable to create {video2x_cache_directory}\nAborting...')
|
||||
raise e
|
||||
else:
|
||||
raise FileNotFoundError('Could not create cache directory')
|
||||
|
||||
# load more settings from gui
|
||||
width = self.width.get()
|
||||
height = self.height.get()
|
||||
scale_ratio = self.scale_ratio.get()
|
||||
image_format = self.image_format.get()
|
||||
threads = self.threads.get()
|
||||
method = AVAILABLE_METHODS[self.method.get()]
|
||||
preserve_frames = self.preserve_frames.get()
|
||||
|
||||
self.upscaler = Upscaler(input_video=input_file, output_video=output_file, method=method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
|
||||
|
||||
# set optional options
|
||||
self.upscaler.waifu2x_driver = driver
|
||||
self.upscaler.scale_width = width
|
||||
self.upscaler.scale_height = height
|
||||
self.upscaler.scale_ratio = scale_ratio
|
||||
self.upscaler.model_dir = None
|
||||
self.upscaler.threads = threads
|
||||
self.upscaler.video2x_cache_directory = video2x_cache_directory
|
||||
self.upscaler.image_format = image_format
|
||||
self.upscaler.preserve_frames = preserve_frames
|
||||
|
||||
# run upscaler
|
||||
self.upscaler.create_temp_directories()
|
||||
|
||||
# start progress bar
|
||||
progress_bar = threading.Thread(target=self._progress_bar)
|
||||
progress_bar.start()
|
||||
|
||||
# start upscaling
|
||||
self.upscaler.run()
|
||||
self.upscaler.cleanup_temp_directories()
|
||||
|
||||
# show message when upscaling completes
|
||||
messagebox.showinfo('Info', 'Upscaling Completed')
|
||||
self.progress_bar['value'] = 100
|
||||
self.running = False
|
||||
self.start_button_text.set('Start')
|
||||
|
||||
|
||||
def _progress_bar(self):
|
||||
""" This method prints a progress bar
|
||||
|
||||
This method prints a progress bar by keeping track
|
||||
of the amount of frames in the input directory
|
||||
and the output directory. This is originally
|
||||
suggested by @ArmandBernard.
|
||||
"""
|
||||
# initialize variables early
|
||||
self.upscaler.progress_bar_exit_signal = False
|
||||
self.upscaler.total_frames_upscaled = 0
|
||||
self.upscaler.total_frames = 1
|
||||
|
||||
# initialize progress bar values
|
||||
self.progress_bar['value'] = 0
|
||||
|
||||
while not self.upscaler.progress_bar_exit_signal:
|
||||
self.progress_bar['value'] = int(100 * self.upscaler.total_frames_upscaled / self.upscaler.total_frames)
|
||||
time.sleep(1)
|
||||
|
||||
def _select_input(self):
|
||||
self.input_file.set(askopenfilename(title='Select Input File'))
|
||||
self.output_file.set(f'{self.input_file.get()}_output.mp4')
|
||||
|
||||
def _select_output(self):
|
||||
self.output_file.set(asksaveasfilename(title='Select Output File'))
|
||||
|
||||
|
||||
def read_config(config_file):
|
||||
""" Reads configuration file
|
||||
|
||||
Returns a dictionary read by JSON.
|
||||
"""
|
||||
with open(config_file, 'r') as raw_config:
|
||||
config = json.load(raw_config)
|
||||
return config
|
||||
|
||||
|
||||
def absolutify_paths(config):
|
||||
""" Check to see if paths to binaries are absolute
|
||||
|
||||
This function checks if paths to binary files are absolute.
|
||||
If not, then absolutify the path.
|
||||
|
||||
Arguments:
|
||||
config {dict} -- configuration file dictionary
|
||||
|
||||
Returns:
|
||||
dict -- configuration file dictionary
|
||||
"""
|
||||
current_directory = pathlib.Path(sys.argv[0]).parent.absolute()
|
||||
|
||||
# check waifu2x-caffe path
|
||||
if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
|
||||
config['waifu2x_caffe']['waifu2x_caffe_path'] = current_directory / config['waifu2x_caffe']['waifu2x_caffe_path']
|
||||
|
||||
# check waifu2x-converter-cpp path
|
||||
if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
|
||||
config['waifu2x_converter']['waifu2x_converter_path'] = current_directory / config['waifu2x_converter']['waifu2x_converter_path']
|
||||
|
||||
# check waifu2x_ncnn_vulkan path
|
||||
if not re.match('^[a-z]:', config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'], re.IGNORECASE):
|
||||
config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = current_directory / config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path']
|
||||
|
||||
# check ffmpeg path
|
||||
if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
|
||||
config['ffmpeg']['ffmpeg_path'] = current_directory / config['ffmpeg']['ffmpeg_path']
|
||||
|
||||
# check video2x cache path
|
||||
if config['video2x']['video2x_cache_directory']:
|
||||
if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
|
||||
config['video2x']['video2x_cache_directory'] = current_directory / config['video2x']['video2x_cache_directory']
|
||||
|
||||
return config
|
||||
|
||||
|
||||
video2x_gui = Video2xGui()
|
||||
Loading…
Reference in New Issue
Block a user