How to write depth data (AVDepthData) to photo file from AVCapturePhoto object?

Question

I've implemented solution similar to Apple's example of capturing depth data with iPhone lidar camera. Main code snippets are as follow:

1. Setting depth formats

let device = AVCaptureDevice.default(.builtInLiDARDepthCamera, for: .video, position: .back)!
let vidFormatsWithDepth = device.formats.filter { format in
    format.formatDescription.dimensions.width == 1920 &&
    format.formatDescription.mediaSubType.rawValue == kCVPixelFormatType_420YpCbCr8BiPlanarFullRange &&
    !format.isVideoBinned &&
    !format.supportedDepthDataFormats.isEmpty &&
    format.supportedDepthDataFormats.contains { $0.formatDescription.mediaSubType.rawValue == kCVPixelFormatType_DepthFloat16 }
}

if let format = vidFormatsWithDepth.first {
    let depthFormats = format.supportedDepthDataFormats.filter { $0.formatDescription.mediaSubType.rawValue == kCVPixelFormatType_DepthFloat16 }
    try! device.lockForConfiguration()
    device.activeFormat = format
    device.activeDepthDataFormat = depthFormats.last
    device.unlockForConfiguration()
}

2. Photo output

func setUpPhotoOutput() {
    photoOutput = AVCapturePhotoOutput()
    photoOutput.maxPhotoQualityPrioritization = .quality
    self.captureSession.addOutput(photoOutput)
    photoOutput.isDepthDataDeliveryEnabled = photoOutput.isDepthDataDeliverySupported
}

3. Capturing photo

var format: [String: Any] = [:]
if photoOutput.availablePhotoPixelFormatTypes.contains(kCVPixelFormatType_420YpCbCr8BiPlanarFullRange) {
    format[kCVPixelBufferPixelFormatTypeKey as String] = kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
}
let settings = AVCapturePhotoSettings(format: format)
settings.isDepthDataDeliveryEnabled = photoOutput.isDepthDataDeliveryEnabled
settings.isDepthDataFiltered = photoOutput.isDepthDataDeliveryEnabled
settings.embedsDepthDataInPhoto = photoOutput.isDepthDataDeliveryEnabled
photoOutput.capturePhoto(with: settings , delegate: self)

4. Processing captured photo data

func createPhotoFile(
    photo: AVCapturePhoto
) {
    let customizer = PhotoDataCustomizer()
    var mainImageData = photo.fileDataRepresentation(with: customizer)!
    // note mainImageData should have embeded depth data, but...
    let imageSource = CGImageSourceCreateWithData(mainImageData as CFData, nil)!
    let depthDataDict = CGImageSourceCopyAuxiliaryDataInfoAtIndex(
        imageSource,
        0,
        kCGImageAuxiliaryDataTypeDepth
    )
    let disparityDataDict = CGImageSourceCopyAuxiliaryDataInfoAtIndex(
        imageSource,
        0,
        kCGImageAuxiliaryDataTypeDisparity
    )
    print("depthDataDict", depthDataDict ?? "nil")
    print("disparityDataDict", disparityDataDict ?? "nil")
    // ... both depthDataDict and disparityDataDict come out as nil
}

class PhotoDataCustomizer: NSObject, AVCapturePhotoFileDataRepresentationCustomizer {
    func replacementDepthData(for photo: AVCapturePhoto) -> AVDepthData? {
        let depthData = photo.depthData?.converting(toDepthDataType: kCVPixelFormatType_DepthFloat16)
        return depthData
    }
}

AVCapturePhoto's photo.depthData is present (not nil) and I would expect that it is embedded if settings.embedsDepthDataInPhoto = true, but both variants of depth data (kCGImageAuxiliaryDataTypeDepth, kCGImageAuxiliaryDataTypeDisparity) come out nil from CGImageSource.

How to properly read the depth data from photo file ... or properly write the depth data in the first place?

Answer 1

I was able to save depth data by adding it manually bypassing var mainImageData = photo.fileDataRepresentation(with: customizer)! with the following:

func createPhotoFile(
    photo: AVCapturePhoto
) {
    let cgMainImage = photo.cgImageRepresentation()!
    let mainImageData = NSMutableData()
    let imageDest = CGImageDestinationCreateWithData(mainImageData, imageFormat as CFString, 1, nil)!
    var options: [String: Any] = photo.metadata
    options[kCGImageDestinationLossyCompressionQuality as String] = 0.8
    options[kCGImagePropertyOrientation as String] = NSNumber(value: orientation.rawValue)
    CGImageDestinationAddImage(imageDest, cgMainImage, options as CFDictionary)

    var depthDataVariant: NSString?
    if let depthData = photo.depthData?.converting(toDepthDataType: kCVPixelFormatType_DepthFloat16)
        .applyingExifOrientation(orientation),
        let depthDataDict = depthData.dictionaryRepresentation(forAuxiliaryDataType: &depthDataVariant),
        let depthDataVariant = depthDataVariant
    {
        CGImageDestinationAddAuxiliaryDataInfo(imageDest, depthDataVariant, depthDataDict as CFDictionary)
    }
    CGImageDestinationFinalize(imageDest)

    // final test ...
    let imageSource = CGImageSourceCreateWithData(mainImageData as CFData, nil)!
    let depthDataDict = CGImageSourceCopyAuxiliaryDataInfoAtIndex(
        imageSource,
        0,
        kCGImageAuxiliaryDataTypeDepth
    )
    print("depthDataDict", depthDataDict ?? "nil")
    // ... now the depthDataDict exists!
}