feat: added base for sdr/hdr detection

This commit is contained in:
Hannes
2026-09-18 19:20:09 +02:00
parent 06c7268308
commit 6d0daddea5
+71 -17
View File
@@ -3,6 +3,7 @@ import sys
import os
import subprocess
import json
import argparse
from check_video import check_video_ext, normalize_language, normalize_lang_code
color = True
@@ -124,7 +125,7 @@ def short_codec_name(codec):
return fixed_width(new_codec.upper(), 5)
class video_lines:
def __init__(self, stream):
def __init__(self, stream, hdr_type):
if stream.get("index"):
self.id = stream.get("index")
else:
@@ -185,6 +186,11 @@ class video_lines:
else:
self.field_order = ""
if hdr_type:
self.hdr_type = hdr_type
else:
self.hdr_type = ""
def __str__(self):
string = fixed_width("Video", 6)
if self.id != None:
@@ -197,7 +203,10 @@ class video_lines:
if self.resolution != "x":
string += f" ({self.resolution}"
if self.framerate != "x":
string += f"@{self.framerate})"
string += f"@{self.framerate}"
string += f" {self.hdr_type})"
elif self.hdr_type != "Unknown HDR":
string += f" ({self.hdr_type})"
if self.aspect_ratio:
string += f" [{self.aspect_ratio}]"
if self.pix_fmt and self.color_space:
@@ -399,7 +408,7 @@ def get_media_info(file):
info = json.loads(result.stdout)
except subprocess.CalledProcessError as e:
print(f"Error running ffprobe: {e.stderr}")
return None, [], [], [], None
return None, [], [], [], None, "UNKNOWN"
# Duration
duration = None
@@ -424,12 +433,41 @@ def get_media_info(file):
elif stream_type == "subtitle":
subtitle_streams.append(stream)
# Detect HDR type from the first video stream
hdr_types = []
if video_streams:
for video in video_streams:
color_transfer = video.get("color_transfer")
color_primaries = video.get("color_primaries")
if color_transfer == "smpte2084":
hdr_types.append("HDR10")
elif color_transfer == "arib-std-b67":
hdr_types.append("HLG")
elif (
color_transfer == "bt709"
and color_primaries == "bt709"
):
hdr_types.append("SDR")
else:
hdr_types.append("Unknown HDR")
if duration is not None:
duration = float(duration)
else:
duration = float('nan')
return duration, video_streams, audio_streams, subtitle_streams, title
return (
duration,
video_streams,
audio_streams,
subtitle_streams,
title,
hdr_types
)
def seconds_to_hms(seconds):
if type(seconds) is float:
@@ -453,9 +491,9 @@ class video_file:
self.name = os.path.basename(path) # 25.mkv
self.size = os.path.getsize(path)
self.duration, videos, audios, subs, title = get_media_info(path)
self.duration, videos, audios, subs, title, hdr_types = get_media_info(path)
self.sort_video_info(videos)
self.sort_video_info(videos, hdr_types)
self.sort_audio_info(audios)
self.sort_subs_info(subs)
@@ -464,11 +502,11 @@ class video_file:
self.duration = seconds_to_hms(self.duration)
self.title = title
def sort_video_info(self, videos):
def sort_video_info(self, videos, hdr_types):
self.videos = []
if videos:
for vl in videos:
video_line = video_lines(vl)
for vl, hdr_t in zip(videos, hdr_types):
video_line = video_lines(vl, hdr_t)
self.videos.append(video_line)
def sort_audio_info(self, audios):
@@ -485,6 +523,7 @@ class video_file:
self.subtitles.append(subtitle)
def print(self):
if not mini:
if self.base_tab == "\t":
np(f"{self.base_tab}{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
else:
@@ -499,6 +538,16 @@ class video_file:
for subtitle in self.subtitles:
np(f"\t\t{subtitle}", NORMAL_STYLE)
print()
else:
if self.base_tab == "\t":
np(f"{self.base_tab}{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
else:
np(f"{os.path.dirname(self.path)}/", INFO_STYLE)
np(f"\t{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
if self.title:
np(f"\t\tTitle: \"{self.title}\"", NORMAL_STYLE)
print()
def get_folder_info(files):
@@ -561,24 +610,30 @@ def handle_folders(dirs, all_files):
all_files.append(dir_files)
if __name__ == "__main__":
# print(sys.argv)
parser = argparse.ArgumentParser()
parser.add_argument("-m", "--mini", action="store_true", help="Enable mini mode")
parser.add_argument("paths", nargs="*", help="Files or directories to process")
args = parser.parse_args()
mini = args.mini
# print(args)
subprocess.run(["python", "/home/honney/.bin/tracker.py", "add", "ff"])
file_dir_array = []
if len(sys.argv) == 0:
print("Something went horribly wrong!")
if len(sys.argv) == 1:
# current_dir = os.path.dirname(os.path.realpath(__file__))
# No paths were provided: use the current directory
if not args.paths:
current_dir = os.getcwd()
handle_folders([os.path.abspath(current_dir)], file_dir_array)
else:
files = []
dirs = []
for argv in sys.argv[1:]:
for argv in args.paths:
if os.path.isfile(argv):
if check_video_ext(os.path.splitext(argv)[1]):
files.append(os.path.abspath(argv))
# else:
# np(f"{os.path.abspath(argv)} is not a compatabile Video file", WARN_STYLE)
# np(
# f"{os.path.abspath(argv)} is not a compatible Video file",
# WARN_STYLE
# )
elif os.path.isdir(argv):
dirs.append(os.path.abspath(argv))
else:
@@ -586,7 +641,6 @@ if __name__ == "__main__":
sys.exit()
handle_folders(dirs, file_dir_array)
handle_files(files, file_dir_array)
for element in file_dir_array:
if type(element) == list:
get_folder_info(element)