feat: added base for sdr/hdr detection

ff

@@ -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:
+            for vl, hdr_t in zip(videos, hdr_types):
-                video_line = video_lines(vl)
+                video_line = video_lines(vl, hdr_t)
                 self.videos.append(video_line)
 
     def sort_audio_info(self, audios):
@@ -485,20 +523,31 @@ class video_file:
                 self.subtitles.append(subtitle)
 
     def print(self):
-        if self.base_tab == "\t":
+        if not mini:
-            np(f"{self.base_tab}{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
+            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)
+            for video in self.videos:
+                np(f"\t\t{video}", NORMAL_STYLE)
+            for audio in self.audios:
+                np(f"\t\t{audio}", NORMAL_STYLE)
+            for subtitle in self.subtitles:
+                np(f"\t\t{subtitle}", NORMAL_STYLE)
+            print()
         else:
-            np(f"{os.path.dirname(self.path)}/", INFO_STYLE)
+            if self.base_tab == "\t":
-            np(f"\t{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
+                np(f"{self.base_tab}{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
-        if self.title:
+            else:
-            np(f"\t\tTitle: \"{self.title}\"", NORMAL_STYLE)
+                np(f"{os.path.dirname(self.path)}/", INFO_STYLE)
-        for video in self.videos:
+                np(f"\t{self.name} ({self.size}, {self.duration}, {self.bitrate} MB/s):", NORMAL_STYLE)
-            np(f"\t\t{video}", NORMAL_STYLE)
+            if self.title:
-        for audio in self.audios:
+                np(f"\t\tTitle: \"{self.title}\"", NORMAL_STYLE)
-            np(f"\t\t{audio}", NORMAL_STYLE)
+            print()
-        for subtitle in self.subtitles:
+
-            np(f"\t\t{subtitle}", NORMAL_STYLE)
-        print()
         
 
 def get_folder_info(files):
@@ -561,32 +610,37 @@ 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:
+    # No paths were provided: use the current directory
-        print("Something went horribly wrong!")
+    if not args.paths:
-    if len(sys.argv) == 1:
-        # current_dir = os.path.dirname(os.path.realpath(__file__))
         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:
-                np(f"This is not a file or directory: {argv}\nNow canceling!", ERROR_STYLE)
+                np(f"This is not a file or directory: {argv}\nNow canceling!",ERROR_STYLE)
                 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)