Merge remote-tracking branch 'origin/master'

bin/upscaler.py

@@ -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}'])
+                    self.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
+                    delta = self.total_frames_upscaled - previous_cycle_frames
-                    previous_cycle_frames = total_frames_upscaled
+                    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

bin/video2x_gui.py

@@ -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()