basic HDR support

jcm93 · Jun 26, 2024 · 1a84cec · 1a84cec
1 parent b83d744
commit 1a84cec
Show file tree

Hide file tree

Showing 8 changed files with 119 additions and 57 deletions.
diff --git a/CaptureSample/Encoder.swift b/CaptureSample/Encoder.swift
@@ -27,6 +27,7 @@ class VTEncoder {
     var decodeSession: VTDecompressionSession!
     var videoSink: VideoSink!
     var pixelTransferSession: VTPixelTransferSession?
+	var hdrMetadataGenerationSession: VTHDRPerFrameMetadataGenerationSession!
     var stoppingEncoding = false
     var pixelTransferBuffer: CVPixelBuffer!
 
@@ -75,6 +76,7 @@ class VTEncoder {
                                        isRealTime: true,
                                        usesReplayBuffer: options.usesReplayBuffer,
                                        replayBufferDuration: options.replayBufferDuration)
+		try self.hdrMetadataGenerationSession = VTHDRPerFrameMetadataGenerationSession(framesPerSecond: 120, hdrFormats: [.dolbyVision])
         if options.convertsColorSpace || options.scales {
             var err2 = VTPixelTransferSessionCreate(allocator: nil, pixelTransferSessionOut: &pixelTransferSession)
             if noErr != err2 {
@@ -96,9 +98,14 @@ class VTEncoder {
     func configureSession(options: Options) async {
         var err: OSStatus = noErr
         if options.codec == kCMVideoCodecType_H264 {
-            err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_H264_Main_AutoLevel)
+			err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_H264_Main_AutoLevel)
         } else if options.codec == kCMVideoCodecType_HEVC {
-            err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_HEVC_Main_AutoLevel)
+			switch options.bitDepth {
+			case 8:
+				err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_HEVC_Main_AutoLevel)
+			default:
+				err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_HEVC_Main10_AutoLevel)
+			}
         }
         if noErr != err {
             logger.fault("Failed to set profile level: \(err, privacy: .public)")
@@ -153,14 +160,18 @@ class VTEncoder {
         if noErr != err {
             logger.fault("Failed to set max keyframe interval duration: \(err, privacy: .public)")
         }
-        err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ColorPrimaries, value: options.colorPrimaries)
-        if noErr != err {
-            logger.fault("Failed to set color primaries: \(err, privacy: .public)")
-        }
-        err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_OutputBitDepth, value: options.bitDepth as CFNumber)
+		err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ColorPrimaries, value: options.colorPrimaries)
+		if noErr != err {
+			logger.fault("Failed to set color primaries: \(err, privacy: .public)")
+		}
+		err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_OutputBitDepth, value: options.bitDepth as CFNumber)
         if noErr != err {
             logger.fault("Failed to set bit depth: \(err, privacy: .public)")
         }
+		err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_HDRMetadataInsertionMode, value: kVTHDRMetadataInsertionMode_Auto)
+		if noErr != err {
+			logger.fault("Failed to set hdr metadata insertion mode: \(err, privacy: .public)")
+		}
         err = VTSessionSetProperty(session, key: kVTCompressionPropertyKey_YCbCrMatrix, value: options.yuvMatrix)
         if noErr != err {
             logger.fault("Failed to set YCbCr matrix: \(err, privacy: .public)")
@@ -207,6 +218,11 @@ class VTEncoder {
         } else {
             pixelBufferToEncodeFrom = buffer
         }
+		do {
+			try self.hdrMetadataGenerationSession.attachMetadata(to: pixelBufferToEncodeFrom)
+		} catch {
+			fatalError()
+		}
         VTCompressionSessionEncodeFrame(self.session, imageBuffer: pixelBufferToEncodeFrom, presentationTimeStamp: timeStamp, duration: duration, frameProperties: properties, infoFlagsOut: infoFlags) {
             (status: OSStatus, infoFlags: VTEncodeInfoFlags, sbuf: CMSampleBuffer?) -> Void in
             if sbuf != nil {
@@ -250,8 +266,8 @@ class VTEncoder {
             }
             /*if let matrix = options.yuvMatrix {
                 CVBufferSetAttachment(buffer, kCVImageBufferYCbCrMatrixKey, matrix, .shouldPropagate)
-            }*/
-            /*if let tf = options.transferFunction {
+            }
+            if let tf = options.transferFunction {
                 CVBufferSetAttachment(buffer, kCVImageBufferTransferFunctionKey, tf, .shouldPropagate)
             }*/
             if pixelTransferSession != nil {
@@ -263,6 +279,11 @@ class VTEncoder {
             } else {
                 pixelBufferToEncodeFrom = buffer
             }
+			do {
+				try self.hdrMetadataGenerationSession.attachMetadata(to: pixelBufferToEncodeFrom)
+			} catch {
+				fatalError()
+			}
             VTCompressionSessionEncodeFrame(self.session, imageBuffer: pixelBufferToEncodeFrom, presentationTimeStamp: timeStamp, duration: duration, frameProperties: properties, infoFlagsOut: infoFlags) {
                 (status: OSStatus, infoFlags: VTEncodeInfoFlags, sbuf: CMSampleBuffer?) -> Void in
                 if sbuf != nil {

diff --git a/CaptureSample/Enums.swift b/CaptureSample/Enums.swift
@@ -1,11 +1,29 @@
 import Foundation
 import VideoToolbox
+import ScreenCaptureKit
 
 public enum CaptureType: Int, Codable, CaseIterable {
     case display
     case window
 }
 
+public enum CaptureHDRStatus: Int, Codable, CaseIterable {
+	case SDR
+	case localHDR
+	case canonicalHDR
+	func enumValue() -> SCCaptureDynamicRange {
+		switch self {
+		case .SDR:
+			.SDR
+		case .localHDR:
+			.hdrLocalDisplay
+		case .canonicalHDR:
+			.hdrCanonicalDisplay
+		}
+	}
+}
+
+
 public enum EncoderSetting: Int, Codable, CaseIterable {
     case H264
     case H265
@@ -185,6 +203,7 @@ public enum CapturePixelFormat: Int, Codable, CaseIterable {
     case l10r
     case biplanarpartial420v
     case biplanarfull420f
+	case biplanarfull444f
     func osTypeFormat() -> OSType {
         switch self {
         case .bgra:
@@ -195,6 +214,8 @@ public enum CapturePixelFormat: Int, Codable, CaseIterable {
             return kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
         case .biplanarfull420f:
             return kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
+		case .biplanarfull444f:
+			return kCVPixelFormatType_444YpCbCr10BiPlanarFullRange
         }
     }
     func stringValue() -> String {
@@ -207,6 +228,8 @@ public enum CapturePixelFormat: Int, Codable, CaseIterable {
             return "420v"
         case .biplanarfull420f:
             return "420f"
+		case .biplanarfull444f:
+			return "xf44"
         }
     }
 }
@@ -215,6 +238,7 @@ public enum CaptureYUVMatrix: Int, Codable, CaseIterable {
     case itu_r_709
     case itu_r_601
     case smpte_240m_1995
+	case none
     func cfStringFormat() -> CFString {
         switch self {
         case .itu_r_709:
@@ -223,6 +247,8 @@ public enum CaptureYUVMatrix: Int, Codable, CaseIterable {
             return CGDisplayStream.yCbCrMatrix_ITU_R_601_4
         case .smpte_240m_1995:
             return CGDisplayStream.yCbCrMatrix_SMPTE_240M_1995
+		case .none:
+			return "" as CFString
         }
     }
     func stringValue() -> String {
@@ -233,6 +259,8 @@ public enum CaptureYUVMatrix: Int, Codable, CaseIterable {
             return "601"
         case .smpte_240m_1995:
             return "SMPTE 240M 1995"
+		case .none:
+			return ""
         }
     }
 }

diff --git a/CaptureSample/Logging.swift b/CaptureSample/Logging.swift
@@ -28,7 +28,8 @@ extension Logger {
         do {
             let logStore = try OSLogStore(scope: .currentProcessIdentifier)
             let timeIntervalToFetch = Date().timeIntervalSince(startupTime)
-            let predicate = NSPredicate(format: "subsystem CONTAINS[c] 'com.jcm.record'")
+            guard let bundleIdentifier = Bundle.main.bundleIdentifier else { return nil }
+            let predicate = NSPredicate(format: "subsystem CONTAINS[c] '\(bundleIdentifier)'")
             let entries = try logStore.getEntries(at: logStore.position(timeIntervalSinceEnd: timeIntervalToFetch), matching: predicate)
             var logString = ""
             for entry in entries {

diff --git a/CaptureSample/ScreenRecorder.swift b/CaptureSample/ScreenRecorder.swift
@@ -198,6 +198,10 @@ class ScreenRecorder: ObservableObject {
     @Published var captureType: CaptureType = .display {
         didSet { updateEngine() }
     }
+
+	@Published var captureHDRStatus: CaptureHDRStatus = .localHDR {
+		didSet { updateEngine() }
+	}
 
     @Published var selectedDisplay: SCDisplay? {
         didSet { updateEngine() }
@@ -406,6 +410,17 @@ class ScreenRecorder: ObservableObject {
         await self.refreshAvailableContent()
     }
 
+    func initializeEventTap() async {
+        if self.eventTap == nil {
+            do {
+                self.eventTap = try RecordEventTap()
+            } catch {
+                logger.fault("Hotkey listener was not initialized: \(error, privacy: .public)")
+            }
+        }
+        self.eventTap?.callback = self.saveReplayBuffer
+    }
+
     /// Starts capturing screen content.
     func start() async {
         // Exit early if already running.
@@ -444,7 +459,7 @@ class ScreenRecorder: ObservableObject {
             // Unable to start the stream. Set the running state to false.
             isRunning = false
         }
-        self.captureEngine.streamOutput.errorHandler = handleEncoderError
+		self.captureEngine.streamOutput.errorHandler = handleEncoderError
     }
 
     /// Stops capturing screen content.
@@ -547,17 +562,6 @@ class ScreenRecorder: ObservableObject {
             await captureEngine.update(configuration: streamConfiguration, filter: contentFilter)
         }
         self.selectedPreset = nil
-        if self.eventTap == nil {
-            do {
-                self.eventTap = try RecordEventTap()
-            } catch {
-                logger.fault("Hotkey listener was not initialized: \(error, privacy: .public)")
-            }
-        }
-        self.eventTap?.callback = self.saveReplayBuffer
-        if self.showsEncodePreview {
-            self.updateEncodePreview()
-        }
     }
 
     func uninstallExtension() {
@@ -614,7 +618,8 @@ class ScreenRecorder: ObservableObject {
 
     private var streamConfiguration: SCStreamConfiguration {
 
-        let streamConfig = SCStreamConfiguration()
+		let streamConfig = SCStreamConfiguration(preset: .captureHDRStreamCanonicalDisplay)
+		streamConfig.captureDynamicRange = captureHDRStatus.enumValue()
 
         // Configure audio capture.
         streamConfig.capturesAudio = isAudioCaptureEnabled
@@ -653,23 +658,18 @@ class ScreenRecorder: ObservableObject {
         // the memory footprint of WindowServer.
         streamConfig.queueDepth = 15
         streamConfig.backgroundColor = CGColor.clear
+		//streamConfig.colorMatrix = "" as CFString
 
         return streamConfig
     }
 
     func assignPixelFormatAndColorMatrix(_ config: SCStreamConfiguration) {
         config.pixelFormat = self.capturePixelFormat.osTypeFormat()
-        if self.bitDepthSetting == .eight {
-            if config.pixelFormat == kCVPixelFormatType_420YpCbCr10BiPlanarFullRange {
-                config.pixelFormat = kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
-            } else if config.pixelFormat == kCVPixelFormatType_420YpCbCr10BiPlanarVideoRange {
-                config.pixelFormat = kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
-            }
-        }
-        if (self.capturePixelFormat == .biplanarpartial420v || self.capturePixelFormat == .biplanarpartial420v) {
+        /*if (self.capturePixelFormat == .biplanarpartial420v || self.capturePixelFormat == .biplanarpartial420v) {
             config.colorMatrix = self.captureYUVMatrix.cfStringFormat()
-        }
+        }*/
         config.colorSpaceName = self.captureColorSpace.cfString()
+		config.colorMatrix = self.captureYUVMatrix.cfStringFormat()
     }
 
     /// - Tag: GetAvailableContent

diff --git a/CaptureSample/Views/CapturePreview.swift b/CaptureSample/Views/CapturePreview.swift
@@ -17,6 +17,7 @@ struct CaptureSingleViewPreview: NSViewRepresentable {
     init() {
         //contentLayer.contentsGravity = .resizeAspect
         contentLayer.contentsGravity = .resizeAspect
+		contentLayer.wantsExtendedDynamicRangeContent = true
     }
 
     func makeNSView(context: Context) -> CaptureVideoPreview {

diff --git a/CaptureSample/Views/Configuration View/VideoCaptureConfigurationView.swift b/CaptureSample/Views/Configuration View/VideoCaptureConfigurationView.swift
@@ -72,6 +72,22 @@ struct VideoCaptureConfigurationView: View {
                     }
                     .labelsHidden()
                     Group {
+						HStack {
+							Text("Capture HDR Status:")
+							Picker("Capture", selection: $screenRecorder.captureHDRStatus) {
+								Text("SDR")
+									.tag(CaptureHDRStatus.SDR)
+								Text("Local HDR")
+									.tag(CaptureHDRStatus.localHDR)
+								Text("Canonical HDR")
+									.tag(CaptureHDRStatus.canonicalHDR)
+							}
+							.pickerStyle(.radioGroup)
+							.horizontalRadioGroupLayout()
+							.alignmentGuide(.imageTitleAlignmentGuide) { dimension in
+								dimension[.leading]
+							}
+						}
                         HStack {
                             Text("Pixel Format:")
                             Picker("Pixel Format", selection: $screenRecorder.capturePixelFormat) {
@@ -85,21 +101,19 @@ struct VideoCaptureConfigurationView: View {
                             }
                             .frame(width: 150)
                         }
-                        if (self.screenRecorder.capturePixelFormat == .biplanarfull420f || self.screenRecorder.capturePixelFormat == .biplanarpartial420v) {
-                            HStack {
-                                Text("Transfer Function:")
-                                Picker("Transfer Function", selection: $screenRecorder.captureYUVMatrix) {
-                                    ForEach(CaptureYUVMatrix.allCases, id: \.self) { format in
-                                        Text(format.stringValue())
-                                            .tag(format)
-                                    }
-                                }
-                                .alignmentGuide(.imageTitleAlignmentGuide) { dimension in
-                                    dimension[.leading]
-                                }
-                                .frame(width: 150)
-                            }
-                        }
+						HStack {
+							Text("Transfer Function:")
+							Picker("Transfer Function", selection: $screenRecorder.captureYUVMatrix) {
+								ForEach(CaptureYUVMatrix.allCases, id: \.self) { format in
+									Text(format.stringValue())
+										.tag(format)
+								}
+							}
+							.alignmentGuide(.imageTitleAlignmentGuide) { dimension in
+								dimension[.leading]
+							}
+							.frame(width: 150)
+						}
                         HStack {
                             Text("Color Space:")
                             Picker("Color Space", selection: $screenRecorder.captureColorSpace) {

diff --git a/CaptureSample/Views/ContentView.swift b/CaptureSample/Views/ContentView.swift
@@ -110,6 +110,7 @@ struct ContentView: View {
         .onAppear {
             Task {
                 if await screenRecorder.canRecord {
+                    await screenRecorder.initializeEventTap()
                     await screenRecorder.start()
                 } else {
                     isUnauthorized = true
@@ -119,9 +120,3 @@ struct ContentView: View {
         }
     }
 }
-
-struct ContentView_Previews: PreviewProvider {
-    static var previews: some View {
-        ContentView()
-    }
-}