AVAssetReaderTrackOutput read HDR frame from a video file.

Hello, I am trying to read video frames using AVAssetReaderTrackOutput. Here is the sample code:

//prepare assets
let asset = AVURLAsset(url: some_url)
let assetReader = try AVAssetReader(asset: asset)

guard let videoTrack = try await asset.loadTracks(withMediaCharacteristic: .visual).first else {
  throw SomeErrorCode.error
}
var readerSettings: [String: Any] = [
  kCVPixelBufferIOSurfacePropertiesKey as String: [String: String]()
]

//check if HDR video
var isHDRDetected: Bool = false
let hdrTracks = try await asset.loadTracks(withMediaCharacteristic: .containsHDRVideo)
if hdrTracks.count > 0 {
  readerSettings[AVVideoAllowWideColorKey as String] = true
  readerSettings[kCVPixelBufferPixelFormatTypeKey as String] =
    kCVPixelFormatType_420YpCbCr10BiPlanarFullRange
  isHDRDetected = true
}

//add output to assetReader
let output = AVAssetReaderTrackOutput(track: videoTrack, outputSettings: readerSettings)

guard assetReader.canAdd(output) else {
  throw SomeErrorCode.error
}
assetReader.add(output)

guard assetReader.startReading() else {
  throw SomeErrorCode.error
}

//add writer ouput settings
let videoOutputSettings: [String: Any] = [
  AVVideoCodecKey: AVVideoCodecType.hevc,
  AVVideoWidthKey: 1920,
  AVVideoHeightKey: 1080,
]

let finalPath = "//some URL oath"

let assetWriter = try AVAssetWriter(outputURL: finalPath, fileType: AVFileType.mov)

guard assetWriter.canApply(outputSettings: videoOutputSettings, forMediaType: AVMediaType.video)
else {
  throw SomeErrorCode.error
}

let assetWriterInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoOutputSettings)

let sourcePixelAttributes: [String: Any] = [
  kCVPixelBufferPixelFormatTypeKey as String: isHDRDetected
    ? kCVPixelFormatType_420YpCbCr10BiPlanarFullRange : kCVPixelFormatType_32ARGB,
  kCVPixelBufferWidthKey as String: 1920,
  kCVPixelBufferHeightKey as String: 1080,
]

//create assetAdoptor
let assetAdaptor = AVAssetWriterInputTaggedPixelBufferGroupAdaptor(
  assetWriterInput: assetWriterInput, sourcePixelBufferAttributes: sourcePixelAttributes)

guard assetWriter.canAdd(assetWriterInput) else {
  throw SomeErrorCode.error
}
assetWriter.add(assetWriterInput)

guard assetWriter.startWriting() else {
  throw SomeErrorCode.error
}
assetWriter.startSession(atSourceTime: CMTime.zero)

//prepare tranfer session
var session: VTPixelTransferSession? = nil
guard
  VTPixelTransferSessionCreate(allocator: kCFAllocatorDefault, pixelTransferSessionOut: &session)
    == noErr, let session
else {
  throw SomeErrorCode.error
}
guard let pixelBufferPool = assetAdaptor.pixelBufferPool else {
  throw SomeErrorCode.error
}

//read through frames
while let nextSampleBuffer = output.copyNextSampleBuffer() {

  autoreleasepool {
    guard let imageBuffer = CMSampleBufferGetImageBuffer(nextSampleBuffer) else {
      return
    }

    //this part copied from (https://developer.apple.com/videos/play/wwdc2023/10181) at 23:58 timestamp
    let attachment = [
      kCVImageBufferYCbCrMatrixKey: kCVImageBufferYCbCrMatrix_ITU_R_2020,
      kCVImageBufferColorPrimariesKey: kCVImageBufferColorPrimaries_ITU_R_2020,
      kCVImageBufferTransferFunctionKey: kCVImageBufferTransferFunction_SMPTE_ST_2084_PQ,
    ]
    CVBufferSetAttachments(imageBuffer, attachment as CFDictionary, .shouldPropagate)

    //now convert to CIImage with HDR data
    let image = CIImage(cvPixelBuffer: imageBuffer)
    let cropped = ""  //here perform some actions like cropping, flipping, etc. and preserve this changes by converting the extent to CGImage first:

    //this part copied from (https://developer.apple.com/videos/play/wwdc2023/10181) at 24:30 timestamp
    guard
      let cgImage = context.createCGImage(
        cropped, from: cropped.extent, format: .RGBA16,
        colorSpace: CGColorSpace(name: CGColorSpace.itur_2100_PQ)!)
    else {
      continue
    }
    //finally convert it back to CIImage
    let newScaledImage = CIImage(cgImage: cgImage)
    //now write it to a new pixelBuffer
    let pixelBufferAttributes: [String: Any] = [
      kCVPixelBufferCGImageCompatibilityKey as String: true,
      kCVPixelBufferCGBitmapContextCompatibilityKey as String: true,
    ]

    var pixelBuffer: CVPixelBuffer?
    CVPixelBufferCreate(
      kCFAllocatorDefault, Int(newScaledImage.extent.width), Int(newScaledImage.extent.height),
      kCVPixelFormatType_420YpCbCr10BiPlanarFullRange, pixelBufferAttributes as CFDictionary,
      &pixelBuffer)

    guard let pixelBuffer else {
      continue
    }

    context.render(newScaledImage, to: pixelBuffer) //context is a CIContext reference

    var pixelTransferBuffer: CVPixelBuffer?
    CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pixelBufferPool, &pixelTransferBuffer)
    guard let pixelTransferBuffer else {
      continue
    }

    // Transfer the image to the pixel buffer.
    guard
      VTPixelTransferSessionTransferImage(session, from: pixelBuffer, to: pixelTransferBuffer)
        == noErr
    else {
      continue
    }

    //finally append to taggedBuffer
  }
}

assetWriterInput.markAsFinished()
await assetWriter.finishWriting()

The result video is not in correct color as the original video. It turns out too bright. If I play around with attachment values, it can be either too dim or too bright but not exactly proper as the original video. What am I missing in my setup? I did find that kCVPixelFormatType_4444AYpCbCr16 can produce proper video output but then I can't convert it to CIImage and so I can't do the CIImage operations that I need. Mainly cropping and resizing the CIImage