made output extensions in batch processing customizable

src/upscaler.py

@@ -4,7 +4,7 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: June 5, 2020
+Last Modified: June 7, 2020
 
 Description: This file contains the Upscaler class. Each
 instance of the Upscaler class is an upscaler on an image or
@@ -86,7 +86,9 @@ class Upscaler:
         self.scale_ratio = None
         self.processes = 1
         self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
-        self.image_format = 'png'
+        self.extracted_frame_format = 'png'
+        self.image_output_extension = '.png'
+        self.video_output_extension = '.mp4'
         self.preserve_frames = False
 
         # other internal members and signals
@@ -234,7 +236,7 @@ class Upscaler:
         if self.driver == 'waifu2x_caffe':
             if (driver_settings['scale_width'] != 0 and driver_settings['scale_height'] == 0 or
                     driver_settings['scale_width'] == 0 and driver_settings['scale_height'] != 0):
-                Avalon.error('Only one of scale_width and scale_height is specified for waifu2x-caffe')
+                Avalon.error(_('Only one of scale_width and scale_height is specified for waifu2x-caffe'))
                 raise AttributeError('only one of scale_width and scale_height is specified for waifu2x-caffe')
 
             # if scale_width and scale_height are specified, ensure scale_ratio is None
@@ -324,8 +326,8 @@ class Upscaler:
         # images need to be renamed to be recognizable for FFmpeg
         if self.driver == 'waifu2x_converter_cpp':
             for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]:
-                renamed = re.sub(f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.image_format}',
+                renamed = re.sub(f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}',
-                                 f'.{self.image_format}',
+                                 f'.{self.extracted_frame_format}',
                                  str(image.name))
                 (self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
 
@@ -403,49 +405,93 @@ class Upscaler:
         self.driver_object.load_configurations(self)
 
         # initialize FFmpeg object
-        self.ffmpeg_object = Ffmpeg(self.ffmpeg_settings, image_format=self.image_format)
+        self.ffmpeg_object = Ffmpeg(self.ffmpeg_settings, extracted_frame_format=self.extracted_frame_format)
 
         # define processing queue
         self.processing_queue = queue.Queue()
 
         Avalon.info(_('Loading files into processing queue'))
-
-        # if input is a list of files
-        if isinstance(self.input, list):
-
-            Avalon.info(_('Loading files from multiple paths'))
         Avalon.debug_info(_('Input path(s): {}').format(self.input))
 
-            # make output directory if it doesn't exist
+        # make output directory if the input is a list or a directory
+        if isinstance(self.input, list) or self.input.is_dir():
             self.output.mkdir(parents=True, exist_ok=True)
 
-            for input_path in self.input:
+        input_files = []
 
+        # if input is single directory
+        # put it in a list for compability with the following code
+        if not isinstance(self.input, list):
+            input_paths = [self.input]
+        else:
+            input_paths = self.input
+
+        # flatten directories into file paths
+        for input_path in input_paths:
+
+            # if the input path is a single file
+            # add the file's path object to input_files
             if input_path.is_file():
-                    output_path = self.output / input_path.name
+                input_files.append(input_path)
-                    self.processing_queue.put((input_path.absolute(), output_path.absolute()))
 
+            # if the input path is a directory
+            # add all files under the directory into the input_files (non-recursive)
             elif input_path.is_dir():
-                    for input_path in [f for f in input_path.iterdir() if f.is_file()]:
+                input_files.extend([f for f in input_path.iterdir() if f.is_file()])
-                        output_path = self.output / input_path.name
-                        self.processing_queue.put((input_path.absolute(), output_path.absolute()))
 
-        # if input specified is single file
+        output_paths = []
-        elif self.input.is_file():
-            Avalon.info(_('Loading single file'))
-            Avalon.debug_info(_('Input path(s): {}').format(self.input))
-            self.processing_queue.put((self.input.absolute(), self.output.absolute()))
 
-        # if input specified is a directory
+        for input_path in input_files:
-        elif self.input.is_dir():
 
-            Avalon.info(_('Loading files from directory'))
+            # get file type
-            Avalon.debug_info(_('Input path(s): {}').format(self.input))
+            # try python-magic if it's available
-            # make output directory if it doesn't exist
+            try:
-            self.output.mkdir(parents=True, exist_ok=True)
+                input_file_mime_type = magic.from_file(str(input_path.absolute()), mime=True)
-            for input_path in [f for f in self.input.iterdir() if f.is_file()]:
+                input_file_type = input_file_mime_type.split('/')[0]
-                output_path = self.output / input_path.name
+                input_file_subtype = input_file_mime_type.split('/')[1]
-                self.processing_queue.put((input_path.absolute(), output_path.absolute()))
+            except Exception:
+                input_file_type = input_file_subtype = None
+
+            # in case python-magic fails to detect file type
+            # try guessing file mime type with mimetypes
+            if input_file_type not in ['image', 'video']:
+                input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
+                input_file_type = input_file_mime_type.split('/')[0]
+                input_file_subtype = input_file_mime_type.split('/')[1]
+
+            # set default output file suffixes
+            # if image type is GIF, default output suffix is also .gif
+            if input_file_mime_type == 'image/gif':
+                output_path = self.output / (input_path.stem + '.gif')
+
+            elif input_file_type == 'image':
+                output_path = self.output / (input_path.stem + self.image_output_extension)
+
+            elif input_file_type == 'video':
+                output_path = self.output / (input_path.stem + self.video_output_extension)
+
+            # if file is none of: image, image/gif, video
+            # skip to the next task
+            else:
+                Avalon.error(_('File {} ({}) neither an image nor a video').format(input_path, input_file_mime_type))
+                Avalon.warning(_('Skipping this file'))
+                continue
+
+            # if there is only one input file
+            # do not modify output file suffix
+            if isinstance(self.input, pathlib.Path) and self.input.is_file():
+                output_path = self.output
+
+            output_path_id = 0
+            while str(output_path) in output_paths:
+                output_path = output_path.parent / pathlib.Path(f'{output_path.stem}_{output_path_id}{output_path.suffix}')
+                output_path_id += 1
+
+            # record output path
+            output_paths.append(str(output_path))
+
+            # push file information into processing queue
+            self.processing_queue.put((input_path.absolute(), output_path.absolute(), input_file_mime_type, input_file_type, input_file_subtype))
 
         # check argument sanity before running
         self._check_arguments()
@@ -461,27 +507,8 @@ class Upscaler:
         try:
             while not self.processing_queue.empty():
 
-                # reset current processing progress for new job
-                self.total_frames_upscaled = 0
-                self.total_frames = 0
-
                 # get new job from queue
-                self.current_input_file, output_path = self.processing_queue.get()
+                self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype = self.processing_queue.get()
-
-                # get file type
-                try:
-                    input_file_mime_type = magic.from_file(str(self.current_input_file.absolute()), mime=True)
-                    input_file_type = input_file_mime_type.split('/')[0]
-                    input_file_subtype = input_file_mime_type.split('/')[1]
-                except Exception:
-                    input_file_type = input_file_subtype = None
-
-                # in case python-magic fails to detect file type
-                # try guessing file mime type with mimetypes
-                if input_file_type not in ['image', 'video']:
-                    input_file_mime_type = mimetypes.guess_type(self.current_input_file.name)[0]
-                    input_file_type = input_file_mime_type.split('/')[0]
-                    input_file_subtype = input_file_mime_type.split('/')[1]
 
                 # start handling input
                 # if input file is a static image
@@ -553,15 +580,6 @@ class Upscaler:
                     self._upscale_frames()
                     Avalon.info(_('Upscaling completed'))
 
-                # if file is none of: image, image/gif, video
-                # skip to the next task
-                else:
-                    Avalon.error(_('File {} ({}) neither an image nor a video').format(self.current_input_file, input_file_mime_type))
-                    Avalon.warning(_('Skipping this file'))
-                    self.processing_queue.task_done()
-                    self.total_processed += 1
-                    continue
-
                 # start handling output
                 # output can be either GIF or video
 
@@ -569,7 +587,7 @@ class Upscaler:
                 if output_path.suffix.lower() == '.gif':
                     Avalon.info(_('Converting extracted frames into GIF image'))
                     gifski_object = Gifski(self.gifski_settings)
-                    self.process_pool.append(gifski_object.make_gif(self.upscaled_frames, output_path, framerate, self.image_format))
+                    self.process_pool.append(gifski_object.make_gif(self.upscaled_frames, output_path, framerate, self.extracted_frame_format))
                     self._wait()
                     Avalon.info(_('Conversion completed'))
 

src/video2x.py

@@ -13,7 +13,7 @@ __      __  _       _                  ___   __   __
 Name: Video2X Controller
 Creator: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: June 4, 2020
+Last Modified: June 7, 2020
 
 Editor: BrianPetkovsek
 Last Modified: June 17, 2019
@@ -206,7 +206,9 @@ gifski_settings = config['gifski']
 gifski_settings['gifski_path'] = os.path.expandvars(gifski_settings['gifski_path'])
 
 # load video2x settings
-image_format = config['video2x']['image_format'].lower()
+extracted_frame_format = config['video2x']['extracted_frame_format'].lower()
+image_output_extension = config['video2x']['image_output_extension']
+video_output_extension = config['video2x']['video_output_extension']
 preserve_frames = config['video2x']['preserve_frames']
 
 # if preserve frames specified in command line
@@ -245,7 +247,9 @@ try:
     upscaler.scale_ratio = video2x_args.ratio
     upscaler.processes = video2x_args.processes
     upscaler.video2x_cache_directory = video2x_cache_directory
-    upscaler.image_format = image_format
+    upscaler.extracted_frame_format = extracted_frame_format
+    upscaler.image_output_extension = image_output_extension
+    upscaler.video_output_extension = video_output_extension
     upscaler.preserve_frames = preserve_frames
 
     # run upscaler

src/video2x.yaml

@@ -169,5 +169,7 @@ gifski:
   quiet: false # Do not show a progress bar
 video2x:
   video2x_cache_directory: null # default: %TEMP%\video2x, directory where cache files are stored, will be deleted if preserve_frames is not set to true
-  image_format: png # png/jpg intermediate file format used for extracted frames during video processing
+  extracted_frame_format: png # png/jpg intermediate file format used for extracted frames during video processing
+  image_output_extension: .png # image output extension during batch processing
+  video_output_extension: .mp4 # video output extension during batch processing
   preserve_frames: false # if set to true, the cache directory won't be cleaned upon task completion

src/video2x_gui.py

@@ -4,7 +4,7 @@
 Creator: Video2X GUI
 Author: K4YT3X
 Date Created: May 5, 2020
-Last Modified: June 5, 2020
+Last Modified: June 7, 2020
 """
 
 # local imports
@@ -17,6 +17,7 @@ from wrappers.ffmpeg import Ffmpeg
 import contextlib
 import datetime
 import json
+import math
 import mimetypes
 import os
 import pathlib
@@ -277,6 +278,8 @@ class Video2XMainWindow(QMainWindow):
         self.driver_combo_box.currentTextChanged.connect(self.update_gui_for_driver)
         self.processes_spin_box = self.findChild(QSpinBox, 'processesSpinBox')
         self.scale_ratio_double_spin_box = self.findChild(QDoubleSpinBox, 'scaleRatioDoubleSpinBox')
+        self.image_output_extension_line_edit = self.findChild(QLineEdit, 'imageOutputExtensionLineEdit')
+        self.video_output_extension_line_edit = self.findChild(QLineEdit, 'videoOutputExtensionLineEdit')
         self.preserve_frames_check_box = self.findChild(QCheckBox, 'preserveFramesCheckBox')
 
         # frame preview
@@ -472,6 +475,9 @@ class Video2XMainWindow(QMainWindow):
             self.config['video2x']['video2x_cache_directory'] = str((pathlib.Path(tempfile.gettempdir()) / 'video2x').absolute())
         self.cache_line_edit.setText(self.config['video2x']['video2x_cache_directory'])
 
+        self.image_output_extension_line_edit.setText(self.config['video2x']['image_output_extension'])
+        self.video_output_extension_line_edit.setText(self.config['video2x']['video_output_extension'])
+
         # load preserve frames settings
         self.preserve_frames_check_box.setChecked(self.config['video2x']['preserve_frames'])
         self.start_button.setEnabled(True)
@@ -926,11 +932,11 @@ class Video2XMainWindow(QMainWindow):
 
                     # otherwise, use .png by default for all images
                     else:
-                        suffix = '.png'
+                        suffix = self.image_output_extension_line_edit.text()
 
                 # if input is video, use .mp4 as output by default
                 elif input_file_type == 'video':
-                    suffix = '.mp4'
+                    suffix = self.video_output_extension_line_edit.text()
 
                 # if failed to detect file type
                 # use input file's suffix
@@ -942,9 +948,9 @@ class Video2XMainWindow(QMainWindow):
             elif input_path.is_dir():
                 output_path = input_path.parent / f'{input_path.stem}_output'
 
-            # try up to 1000 times
+            # try a new name with a different file ID
             output_path_id = 0
-            while output_path.exists() and output_path_id <= 1000:
+            while output_path.exists():
                 if input_path.is_file():
                     output_path = input_path.parent / pathlib.Path(f'{input_path.stem}_output_{output_path_id}{suffix}')
                 elif input_path.is_dir():
@@ -1067,7 +1073,7 @@ It\'s also highly recommended for you to attach the [log file]({}) under the pro
 
         # upscale process will stop at 99%
         # so it's set to 100 manually when all is done
-        progress_callback.emit((upscale_begin_time, 0, 0, 0, 0, pathlib.Path(), pathlib.Path()))
+        # progress_callback.emit((upscale_begin_time, 0, 0, 0, 0, pathlib.Path(), pathlib.Path()))
 
     def set_progress(self, progress_information: tuple):
         upscale_begin_time = progress_information[0]
@@ -1177,7 +1183,9 @@ It\'s also highly recommended for you to attach the [log file]({}) under the pro
             self.upscaler.scale_ratio = self.scale_ratio_double_spin_box.value()
             self.upscaler.processes = self.processes_spin_box.value()
             self.upscaler.video2x_cache_directory = pathlib.Path(os.path.expandvars(self.cache_line_edit.text()))
-            self.upscaler.image_format = self.config['video2x']['image_format'].lower()
+            self.upscaler.extracted_frame_format = self.config['video2x']['extracted_frame_format'].lower()
+            self.upscaler.image_output_extension = self.image_output_extension_line_edit.text()
+            self.upscaler.video_output_extension = self.video_output_extension_line_edit.text()
             self.upscaler.preserve_frames = bool(self.preserve_frames_check_box.isChecked())
 
             # run upscaler
@@ -1200,6 +1208,10 @@ It\'s also highly recommended for you to attach the [log file]({}) under the pro
             self.upscale_errored(e)
 
     def upscale_errored(self, exception: Exception):
+        # send stop signal in case it's not triggered
+        with contextlib.suppress(AttributeError):
+            self.upscaler.running = False
+
         self.show_error(exception)
         self.threadpool.waitForDone(5)
         self.start_button.setEnabled(True)

src/video2x_gui.ui

@@ -7,7 +7,7 @@
     <x>0</x>
     <y>0</y>
     <width>673</width>
-    <height>844</height>
+    <height>876</height>
    </rect>
   </property>
   <property name="acceptDrops">
@@ -407,6 +407,54 @@
                </item>
               </layout>
              </item>
+             <item>
+              <layout class="QHBoxLayout" name="imageOutputExtensionHorizontalLayout">
+               <item>
+                <widget class="QLabel" name="imageOutputExtensionLabel">
+                 <property name="text">
+                  <string>Image Output Extension</string>
+                 </property>
+                </widget>
+               </item>
+               <item>
+                <widget class="QLineEdit" name="imageOutputExtensionLineEdit">
+                 <property name="sizePolicy">
+                  <sizepolicy hsizetype="Preferred" vsizetype="Fixed">
+                   <horstretch>0</horstretch>
+                   <verstretch>0</verstretch>
+                  </sizepolicy>
+                 </property>
+                 <property name="text">
+                  <string>.png</string>
+                 </property>
+                </widget>
+               </item>
+              </layout>
+             </item>
+             <item>
+              <layout class="QHBoxLayout" name="videoOutputExtensionHorizontalLayout">
+               <item>
+                <widget class="QLabel" name="videoOutputExtensionLabel">
+                 <property name="text">
+                  <string>Video Output Extension</string>
+                 </property>
+                </widget>
+               </item>
+               <item>
+                <widget class="QLineEdit" name="videoOutputExtensionLineEdit">
+                 <property name="sizePolicy">
+                  <sizepolicy hsizetype="Preferred" vsizetype="Fixed">
+                   <horstretch>0</horstretch>
+                   <verstretch>0</verstretch>
+                  </sizepolicy>
+                 </property>
+                 <property name="text">
+                  <string>.mp4</string>
+                 </property>
+                </widget>
+               </item>
+              </layout>
+             </item>
              <item>
               <widget class="QCheckBox" name="preserveFramesCheckBox">
                <property name="toolTip">
@@ -448,6 +496,9 @@
                    <property name="text">
                     <string>Keep Aspect Ratio</string>
                    </property>
+                   <property name="checked">
+                    <bool>true</bool>
+                   </property>
                   </widget>
                  </item>
                  <item>