feat: added base for sdr/hdr detection
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user