2.8.1 added automatic pixel format and color depth detection

bin/ffmpeg.py

@@ -6,7 +6,7 @@
 Name: FFMPEG Class
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: June 15, 2019
+Last Modified: July 9, 2019
 
 Description: This class handles all FFMPEG related
 operations.
@@ -33,6 +33,40 @@ class Ffmpeg:
         self.ffmpeg_binary = os.path.join(self.ffmpeg_path, 'ffmpeg.exe')
         self.ffmpeg_probe_binary = os.path.join(self.ffmpeg_path, 'ffprobe.exe')
         self.image_format = image_format
+        self.pixel_format = None
+    
+    def get_pixel_formats(self):
+        """ Get a dictionary of supported pixel formats
+
+        List all supported pixel formats and their corresponding
+        bit depth.
+
+        Returns:
+            dictionary -- JSON dict of all pixel formats to bit depth
+        """
+        execute = [
+            self.ffmpeg_probe_binary,
+            '-v',
+            'quiet',
+            '-pix_fmts'
+        ]
+
+        Avalon.debug_info(f'Executing: {" ".join(execute)}')
+
+        # initialize dictionary to store pixel formats
+        pixel_formats = {}
+
+        # record all pixel formats into dictionary
+        for line in subprocess.run(execute, check=True, stdout=subprocess.PIPE).stdout.decode().split('\n'):
+            try:
+                pixel_formats[' '.join(line.split()).split()[1]] = int(' '.join(line.split()).split()[3])
+            except (IndexError, ValueError):
+                pass
+
+        # print pixel formats for debugging
+        Avalon.debug_info(pixel_formats)
+
+        return pixel_formats
 
     def get_video_info(self, input_video):
         """ Gets input video information
@@ -179,6 +213,13 @@ class Ffmpeg:
         # from only that section
         if section:
             source = self.ffmpeg_settings[phase][section].keys()
+
+            # if pixel format is not specified, use the source pixel format
+            try:
+                if self.ffmpeg_settings[phase][section].get('-pix_fmt') is None:
+                    self.ffmpeg_settings[phase][section]['-pix_fmt'] = self.pixel_format
+            except KeyError:
+                pass
         else:
             source = self.ffmpeg_settings[phase].keys()
 

bin/upscaler.py

@@ -6,7 +6,7 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: June 26, 2019
+Last Modified: July 9, 2019
 
 Dev: SAT3LL
 
@@ -212,7 +212,7 @@ class Upscaler:
 
                 # create a separate w2 instance for each thread
                 if self.waifu2x_driver == 'waifu2x_caffe':
-                    w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
+                    w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir, self.bit_depth)
                     if self.scale_ratio:
                         thread = threading.Thread(target=w2.upscale,
                                                   args=(thread_info[0],
@@ -314,6 +314,17 @@ class Upscaler:
 
         # get average frame rate of video stream
         framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
+        fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']
+
+        # get a dict of all pixel formats and corresponding bit depth
+        pixel_formats = fm.get_pixel_formats()
+
+        try:
+            self.bit_depth = pixel_formats[fm.pixel_format]
+        except KeyError as e:
+            Avalon.error(f'Unsupported pixel format: {fm.pixel_format}')
+            raise e
+
         Avalon.info(f'Framerate: {framerate}')
 
         # width/height will be coded width/height x upscale factor

bin/video2x.json

@@ -66,7 +66,7 @@
         "-vcodec": "libx264",
         "-crf": 17,
         "-b:v": null,
-        "-pix_fmt": "yuv420p"
+        "-pix_fmt": null
       },
       "-y": true
     },
@@ -76,7 +76,7 @@
         "-map": "1?",
         "-c": "copy",
         "-map": "-1:v?",
-        "-pix_fmt": "yuv420p"
+        "-pix_fmt": null
       },
       "-y": true
     }

bin/video2x.py

@@ -15,7 +15,7 @@ __      __  _       _                  ___   __   __
 Name: Video2X Controller
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: June 26, 2019
+Last Modified: July 9, 2019
 
 Dev: BrianPetkovsek
 Dev: SAT3LL
@@ -57,7 +57,7 @@ import tempfile
 import time
 import traceback
 
-VERSION = '2.8.0'
+VERSION = '2.8.1'
 
 # each thread might take up to 2.5 GB during initialization.
 # (system memory, not to be confused with GPU memory)

bin/waifu2x_caffe.py

@@ -6,7 +6,7 @@
 Name: Waifu2x Caffe Driver
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: June 15, 2019
+Last Modified: July 9, 2019
 
 Description: This class is a high-level wrapper
 for waifu2x-caffe.
@@ -25,10 +25,11 @@ class Waifu2xCaffe:
     the upscale function.
     """
 
-    def __init__(self, waifu2x_settings, process, model_dir):
+    def __init__(self, waifu2x_settings, process, model_dir, bit_depth):
         self.waifu2x_settings = waifu2x_settings
         self.waifu2x_settings['process'] = process
         self.waifu2x_settings['model_dir'] = model_dir
+        self.waifu2x_settings['output_depth'] = bit_depth
 
         # arguments passed through command line overwrites config file values
         self.process = process