Add some checks for non video streams and streams without color data.

Author: Nyanraltotlapun

.gitignore

@@ -157,4 +157,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/

get_hdr_metadata.py

@@ -204,16 +204,15 @@ def to_x265_params(self) -> str:
         return ""
 
     # matrix-coefficients=<arg> Matrix coefficients (CICP) of input content:
-    # identity, bt709, unspecified, fcc73, bt470bg, bt601, smpte240, ycgco, bt2020ncl, bt2020cl, smpte2085, chromncl, chromcl, ictcp
+    # identity, bt709, unspecified, fcc73, bt470bg, bt601, smpte240,
+    # ycgco, bt2020ncl, bt2020cl, smpte2085, chromncl, chromcl, ictcp
     def to_libaom_av1_params(self) -> str:
         res = f"color-primaries={self.color_primaries}:transfer-characteristics={self.color_transfer}"
         mc = libaom_get_matrix_coefficients(self.color_space)
         if mc is not None:
             res += f":matrix-coefficients={mc}"
         return res
 
-
-
     # From Fastfix
     # if (fastflix.current_video.color_space and "bt2020" in fastflix.current_video.color_space):
     #             svtav1_params.append(f"matrix-coefficients=9")
@@ -225,6 +224,51 @@ def to_libsvtav1_params(self) -> str:
         return res
 
 
+def parse_frame_data(frame_data: dict):
+    color_params = ["pix_fmt", "color_space", "color_primaries", "color_transfer"]
+
+    missing_params = [x for x in color_params if x not in frame_data.keys()]
+    if len(missing_params) != 0:
+        print(f"Missing {missing_params} parameters in frame metadata!")
+        print("Probably not an HDR stream!")
+        print("Exit!")
+        return
+
+    color_data = ColorData(frame_data)
+    print("Color Data:")
+    print(color_data)
+    print("")
+
+    x265_params: str = color_data.to_x265_params()
+    libaom_av1_params: str = color_data.to_libaom_av1_params()
+    libsvtav1_params: str = color_data.to_libsvtav1_params()
+
+    side_data_list = frame_data["side_data_list"]
+    for side_data in side_data_list:
+        if side_data["side_data_type"] == "Mastering display metadata":
+            mastering_display_data = MasteringDisplayData(side_data)
+            x265_params += ":" + mastering_display_data.to_x265_params()
+            libsvtav1_params += ":" + mastering_display_data.to_libsvtav1_params()
+            print("Mastering display metadata:")
+            print(mastering_display_data)
+            print("")
+
+        elif side_data["side_data_type"] == "Content light level metadata":
+            content_light_level_data = ContentLightLevelData(side_data)
+            x265_params += ":" + content_light_level_data.to_x265_params()
+            libsvtav1_params += ":" + content_light_level_data.to_libsvtav1_params()
+            print("Content light level metadata:")
+            print(content_light_level_data)
+            print("")
+
+    print(f"\nFFmpeg options: {color_data.to_ffmpeg_options()}\n")
+    print(f"x265 params: {x265_params}\n")
+    print(f"libsvtav1 params: {libsvtav1_params}\n")
+    print(f"libaom-av1 params: {libaom_av1_params}\n")
+
+    print("Done!")
+
+
 if __name__ == '__main__':
 
     # Initialize arguments parser
@@ -272,48 +316,23 @@ def to_libsvtav1_params(self) -> str:
                               "-select_streams", str(arguments.input_stream),
                               "-print_format", "json", "-show_frames", "-read_intervals", "%+#1",
                               "-show_entries",
+                              "stream=codec_type:" +
                               "frame=pix_fmt,color_space,color_primaries,color_transfer,side_data_list",
                               "-i", arguments.input_file]
 
     try:
         result = subprocess.run(ffprobe_cmd, capture_output=True, encoding="UTF-8")
         if result.returncode == 0 and result.stdout is not None:
             metadata = json.loads(result.stdout)
-            frame_data = metadata["frames"][0]
 
-            color_data = ColorData(frame_data)
-            print("Color Data:")
-            print(color_data)
-            print("")
+            stream_codec_type = metadata["streams"][0]["codec_type"]
+            if stream_codec_type == "video":
 
-            x265_params: str = color_data.to_x265_params()
-            libaom_av1_params: str = color_data.to_libaom_av1_params()
-            libsvtav1_params: str = color_data.to_libsvtav1_params()
-
-            side_data_list = frame_data["side_data_list"]
-            for side_data in side_data_list:
-                if side_data["side_data_type"] == "Mastering display metadata":
-                    mastering_display_data = MasteringDisplayData(side_data)
-                    x265_params += ":" + mastering_display_data.to_x265_params()
-                    libsvtav1_params += ":" + mastering_display_data.to_libsvtav1_params()
-                    print("Mastering display metadata:")
-                    print(mastering_display_data)
-                    print("")
-
-                elif side_data["side_data_type"] == "Content light level metadata":
-                    content_light_level_data = ContentLightLevelData(side_data)
-                    x265_params += ":" + content_light_level_data.to_x265_params()
-                    libsvtav1_params += ":" + content_light_level_data.to_libsvtav1_params()
-                    print("Content light level metadata:")
-                    print(content_light_level_data)
-                    print("")
-
-            print(f"\nFFmpeg options: {color_data.to_ffmpeg_options()}\n")
-            print(f"x265 params: {x265_params}\n")
-            print(f"libsvtav1 params: {libsvtav1_params}\n")
-            print(f"libaom-av1 params: {libaom_av1_params}\n")
-
-            print("Done!")
+                parse_frame_data(metadata["frames"][0])
+            else:
+                print(f"Selected stream type is \"{stream_codec_type}\"")
+                print("Not a video stream!")
+                print("Exit!")
 
         else:
             print("Error executing ffprobe binary!")