From ca90c5be02e2fc1981be43eaef57bf7dfa3016dc Mon Sep 17 00:00:00 2001 From: k4yt3x Date: Tue, 9 Jul 2019 20:47:55 -0400 Subject: [PATCH] 2.8.1 added automatic pixel format and color depth detection --- bin/ffmpeg.py | 43 ++++++++++++++++++++++++++++++++++++++++++- bin/upscaler.py | 15 +++++++++++++-- bin/video2x.json | 4 ++-- bin/video2x.py | 4 ++-- bin/waifu2x_caffe.py | 5 +++-- 5 files changed, 62 insertions(+), 9 deletions(-) diff --git a/bin/ffmpeg.py b/bin/ffmpeg.py index 2a3fe9d..565730b 100644 --- a/bin/ffmpeg.py +++ b/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() diff --git a/bin/upscaler.py b/bin/upscaler.py index c3b603d..ea31f84 100644 --- a/bin/upscaler.py +++ b/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 diff --git a/bin/video2x.json b/bin/video2x.json index 79178e1..049e3d3 100644 --- a/bin/video2x.json +++ b/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 } diff --git a/bin/video2x.py b/bin/video2x.py index 9c0ccfe..06f2cdd 100644 --- a/bin/video2x.py +++ b/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) diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 37e7f68..0d18592 100644 --- a/bin/waifu2x_caffe.py +++ b/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