fix(mmclient): support decoding HDR with VideoToolbox on mac

colinmarc · Apr 30, 2024 · 1a66a3f · 1a66a3f
1 parent bae094e
commit 1a66a3f
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 26 deletions.
diff --git a/mm-client/src/render.slang b/mm-client/src/render.slang
@@ -78,7 +78,7 @@ float3 bt2020_pq_to_display(float3 color, int vk_color_space)
     switch (vk_color_space)
     {
     case VK_COLOR_SPACE_SRGB_NONLINEAR_EXT:
-        return srgb_inverse_eotf(clamp(linear, 0.0, 1.0));
+        return srgb_inverse_eotf(linear);
     VK_COLOR_SPACE_BT709_NONLINEAR_EXT:
         return bt709_inverse_eotf(clamp(linear, 0.0, 1.0));
     case VK_COLOR_SPACE_EXTENDED_SRGB_LINEAR_EXT:

diff --git a/mm-client/src/video.rs b/mm-client/src/video.rs
@@ -392,8 +392,28 @@ impl DecoderInit {
             Ok(()) => {
                 self.first_frame = match frame.format() {
                     ffmpeg::format::Pixel::VIDEOTOOLBOX => {
+                        let sw_format = unsafe {
+                            let ctx_ref = (*self.decoder.as_ptr()).hw_frames_ctx;
+                            assert!(!ctx_ref.is_null());
+
+                            let mut transfer_fmt_list = std::ptr::null_mut();
+                            if ffmpeg_sys::av_hwframe_transfer_get_formats(
+                            ctx_ref,
+                            ffmpeg_sys::AVHWFrameTransferDirection::AV_HWFRAME_TRANSFER_DIRECTION_FROM,
+                            &mut transfer_fmt_list,
+                            0) < 0
+                                {
+                                    bail!("call to av_hwframe_transfer_get_formats failed");
+                                };
+
+                            let transfer_formats = read_format_list(transfer_fmt_list);
+                            assert!(!transfer_formats.is_empty());
+
+                            transfer_formats[0]
+                        };
+
                         let mut sw_frame = ffmpeg::frame::Video::new(
-                            ffmpeg::format::Pixel::P010LE, // xxx
+                            sw_format,
                             self.decoder.width(),
                             self.decoder.height(),
                         );
@@ -1025,38 +1045,51 @@ fn copy_frame(
 #[no_mangle]
 unsafe extern "C" fn get_hw_format_videotoolbox(
     ctx: *mut ffmpeg_sys::AVCodecContext,
-    mut formats: *const ffmpeg_sys::AVPixelFormat,
+    list: *const ffmpeg_sys::AVPixelFormat,
 ) -> ffmpeg_sys::AVPixelFormat {
     use ffmpeg_sys::AVPixelFormat::*;
 
-    while *formats != AV_PIX_FMT_NONE {
-        if *formats == AV_PIX_FMT_VIDEOTOOLBOX {
-            let frames_ctx_ref = ffmpeg_sys::av_hwframe_ctx_alloc((*ctx).hw_device_ctx);
-            if frames_ctx_ref.is_null() {
-                error!("call to av_hwframe_ctx_alloc failed");
-                break;
-            }
+    let sw_pix_fmt = (*ctx).sw_pix_fmt;
+    let formats = read_format_list(list);
+
+    if formats.contains(&ffmpeg::format::Pixel::VIDEOTOOLBOX) {
+        let frames_ctx_ref = ffmpeg_sys::av_hwframe_ctx_alloc((*ctx).hw_device_ctx);
+        if frames_ctx_ref.is_null() {
+            error!("call to av_hwframe_ctx_alloc failed");
+            return sw_pix_fmt;
+        }
 
-            let frames_ctx = (*frames_ctx_ref).data as *mut ffmpeg_sys::AVHWFramesContext;
-            (*frames_ctx).width = (*ctx).width;
-            (*frames_ctx).height = (*ctx).height;
-            (*frames_ctx).format = AV_PIX_FMT_VIDEOTOOLBOX;
-            (*frames_ctx).sw_format = (*ctx).sw_pix_fmt;
+        debug!(?formats, sw_pix_fmt = ?sw_pix_fmt, "get_hw_format_videotoolbox");
 
-            let res = ffmpeg_sys::av_hwframe_ctx_init(frames_ctx_ref);
-            if res < 0 {
-                error!("call to av_hwframe_ctx_init failed");
-                break;
-            }
+        let frames_ctx = (*frames_ctx_ref).data as *mut ffmpeg_sys::AVHWFramesContext;
+        (*frames_ctx).width = (*ctx).width;
+        (*frames_ctx).height = (*ctx).height;
+        (*frames_ctx).format = AV_PIX_FMT_VIDEOTOOLBOX;
+        (*frames_ctx).sw_format = AV_PIX_FMT_YUV420P;
 
-            debug!("using VideoToolbox hardware encoder");
-            (*ctx).hw_frames_ctx = frames_ctx_ref;
-            return *formats;
+        let res = ffmpeg_sys::av_hwframe_ctx_init(frames_ctx_ref);
+        if res < 0 {
+            error!("call to av_hwframe_ctx_init failed");
+            return sw_pix_fmt;
         }
 
-        formats = formats.add(1);
+        debug!("using VideoToolbox hardware encoder");
+        (*ctx).hw_frames_ctx = frames_ctx_ref;
+        return AV_PIX_FMT_VIDEOTOOLBOX;
+    }
+
+    warn!("unable to determine ffmpeg hw format");
+    sw_pix_fmt
+}
+
+unsafe fn read_format_list(
+    mut ptr: *const ffmpeg_sys::AVPixelFormat,
+) -> Vec<ffmpeg::format::Pixel> {
+    let mut formats = Vec::new();
+    while !ptr.is_null() && *ptr != ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
+        formats.push((*ptr).into());
+        ptr = ptr.add(1);
     }
 
-    warn!("VideoToolbox setup failed, falling back to CPU decoder");
-    AV_PIX_FMT_YUV420P10LE
+    formats
 }