Как преобразовать текстуру RGBA в текстуры Y и CbCr в металле

Question

У Apple есть полезный учебник под названием Отображение опыта AR с металлом , который показывает, как извлечь текстуры Y и CbCr из свойства ARFrame capturedImage и преобразовать их в RGB для рендеринга. ,Однако я столкнулся с проблемами, пытаясь взять текстуру RGBA и выполнить обратную операцию, т.е. преобразовать обратно в текстуры Y и CbCr.

Я переписал фрагментный шейдер в учебнике как вычислительный шейдер, который пишет втекстура rgba, которую я создал из металлического буфера:

// Same as capturedImageFragmentShader but it's a kernel function instead
kernel void yCbCrToRgbKernel(texture2d<float, access::sample> yTexture [[ texture(kTextureIndex_Y) ]],
                             texture2d<float, access::sample> cbCrTexture [[ texture(kTextureIndex_CbCr) ]],
                             texture2d<float, access::write> rgbaTexture [[ texture(kTextureIndex_RGBA) ]],
                             uint2 gid [[ thread_position_in_grid ]])
{
    constexpr sampler colorSampler(mip_filter::linear, mag_filter::linear, min_filter::linear);

    const float4x4 ycbcrToRGBTransform = float4x4(
        float4(+1.0000f, +1.0000f, +1.0000f, +0.0000f),
        float4(+0.0000f, -0.3441f, +1.7720f, +0.0000f),
        float4(+1.4020f, -0.7141f, +0.0000f, +0.0000f),
        float4(-0.7010f, +0.5291f, -0.8860f, +1.0000f)
    );

    float4 ycbcr = float4(yTexture.sample(colorSampler, float2(gid)).r, cbCrTexture.sample(colorSampler, float2(gid)).rg, 1.0);
    float4 result = ycbcrToRGBTransform * ycbcr;
    rgbaTexture.write(result, ushort2(gid));
}

Я попытался написать второй вычислительный шейдер, чтобы выполнить обратную операцию, вычисляя значения Y, Cb и Cr, используя формулы преобразования, найденные в YCbCr страница википедии :

kernel void rgbaToYCbCrKernel(texture2d<float, access::write> yTexture [[ texture(kTextureIndex_Y) ]],
                             texture2d<float, access::write> cbCrTexture [[ texture(kTextureIndex_CbCr) ]],
                             texture2d<float, access::sample> rgbaTexture [[ texture(kTextureIndex_RGBA) ]],
                             uint2 gid [[ thread_position_in_grid ]])
{
    constexpr sampler colorSampler(mip_filter::linear, mag_filter::linear, min_filter::linear);

    float4 rgba = rgbaTexture.sample(colorSampler, float2(gid)).rgba;

    // see https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion for conversion formulae

    float Y = 16.0 + (65.481 * rgba.r + 128.553 * rgba.g + 24.966 * rgba.b);
    float Cb = 128 + (-37.797 * rgba.r + 74.203 * rgba.g + 112.0 * rgba.b);
    float Cr = 128 + (112.0 * rgba.r + 93.786 * rgba.g - 18.214 * rgba.b);

    yTexture.write(Y, gid);
    cbCrTexture.write(float4(Cb, Cr, 0, 0), gid); // this probably is not correct...
}

Моя проблема заключается в том, как правильно записать данные в эти текстуры. Я знаю, что это неправильно, потому что в результате отображается сплошной розовый цвет. Ожидаемый результат - это, очевидно, оригинальный, немодифицированный дисплей.

Пиксельные форматы для текстур Y, CbCr и RGBA: .r8UNorm, .rg8UNorm и rgba8UNorm соответственно.

Вот мой быстрый код для настройки текстур и выполнения шейдеров:

private func createTexture(fromPixelBuffer pixelBuffer: CVPixelBuffer, pixelFormat: MTLPixelFormat, planeIndex: Int) -> MTLTexture? {
        guard CVMetalTextureCacheCreate(kCFAllocatorSystemDefault, nil, device, nil, &capturedImageTextureCache) == kCVReturnSuccess else { return nil }

        var mtlTexture: MTLTexture? = nil
        let width = CVPixelBufferGetWidthOfPlane(pixelBuffer, planeIndex)
        let height = CVPixelBufferGetHeightOfPlane(pixelBuffer, planeIndex)

        var texture: CVMetalTexture? = nil
        let status = CVMetalTextureCacheCreateTextureFromImage(nil, capturedImageTextureCache!, pixelBuffer, nil, pixelFormat, width, height, planeIndex, &texture)
        if status == kCVReturnSuccess {
            mtlTexture = CVMetalTextureGetTexture(texture!)
        }

        return mtlTexture
    }

    func arFrameToRGB(frame: ARFrame) {

        let frameBuffer = frame.capturedImage

        CVPixelBufferLockBaseAddress(frameBuffer, CVPixelBufferLockFlags(rawValue: 0))

        // Extract Y and CbCr textures
        let capturedImageTextureY = createTexture(fromPixelBuffer: frameBuffer, pixelFormat: .r8Unorm, planeIndex: 0)!
        let capturedImageTextureCbCr = createTexture(fromPixelBuffer: frameBuffer, pixelFormat: .rg8Unorm, planeIndex: 1)!

        // create the RGBA texture
        let rgbaBufferWidth = CVPixelBufferGetWidthOfPlane(frameBuffer, 0)
        let rgbaBufferHeight = CVPixelBufferGetHeightOfPlane(frameBuffer, 0)
        if rgbaBuffer == nil {
            rgbaBuffer = device.makeBuffer(length: 4 * rgbaBufferWidth * rgbaBufferHeight, options: [])
        }

        let rgbaTextureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba8Unorm, width: rgbaBufferWidth, height: rgbaBufferHeight, mipmapped: false)
        rgbaTextureDescriptor.usage = [.shaderWrite, .shaderRead]
        let rgbaTexture = rgbaBuffer?.makeTexture(descriptor: rgbaTextureDescriptor, offset: 0, bytesPerRow: 4 * rgbaBufferWidth)

        threadGroupSize = MTLSizeMake(4, 4, 1)
        threadGroupCount = MTLSizeMake((rgbaTexture!.width + threadGroupSize!.width - 1) / threadGroupSize!.width, (rgbaTexture!.height + threadGroupSize!.height - 1) / threadGroupSize!.height, 1)

        let yCbCrToRGBACommandBuffer = commandQueue.makeCommandBuffer()!
        let yCbCrToRGBAComputeEncoder = yCbCrToRGBACommandBuffer.makeComputeCommandEncoder()!
        yCbCrToRGBAComputeEncoder.setComputePipelineState(yCbCrToRgbPso)
        yCbCrToRGBAComputeEncoder.setTexture(capturedImageTextureY, index: Int(kTextureIndex_Y.rawValue))
        yCbCrToRGBAComputeEncoder.setTexture(capturedImageTextureCbCr, index: Int(kTextureIndex_CbCr.rawValue))
        yCbCrToRGBAComputeEncoder.setTexture(rgbaTexture, index: Int(kTextureIndex_RGBA.rawValue))
        yCbCrToRGBAComputeEncoder.dispatchThreadgroups(threadGroupCount!, threadsPerThreadgroup: threadGroupSize!)
        yCbCrToRGBAComputeEncoder.endEncoding()

        let rgbaToYCbCrCommandBuffer = commandQueue.makeCommandBuffer()!
        let rgbaToYCbCrComputeEncoder = rgbaToYCbCrCommandBuffer.makeComputeCommandEncoder()!
        rgbaToYCbCrComputeEncoder.setComputePipelineState(rgbaToYCbCrPso)
        rgbaToYCbCrComputeEncoder.setTexture(capturedImageTextureY, index: Int(kTextureIndex_Y.rawValue))
        rgbaToYCbCrComputeEncoder.setTexture(capturedImageTextureCbCr, index: Int(kTextureIndex_CbCr.rawValue))
        rgbaToYCbCrComputeEncoder.setTexture(rgbaTexture, index: Int(kTextureIndex_RGBA.rawValue))
        rgbaToYCbCrComputeEncoder.dispatchThreadgroups(threadGroupCount!, threadsPerThreadgroup: threadGroupSize!)
        rgbaToYCbCrComputeEncoder.endEncoding()

        yCbCrToRGBACommandBuffer.commit()
        rgbaToYCbCrCommandBuffer.commit()

        yCbCrToRGBACommandBuffer.waitUntilCompleted()
        rgbaToYCbCrCommandBuffer.waitUntilCompleted()

        CVPixelBufferUnlockBaseAddress(frameBuffer, CVPixelBufferLockFlags(rawValue: 0))
    }

Конечная цель - использовать металлические шейдеры для обработки изображений на текстуре rgba и, в конечном итоге, для обратной записи в текстуры Y и CbCr. для отображения на экране.

Вот части, в которых я не уверен

1. Как мне записать данные в правильном формате в эти текстуры, учитывая, что тип длятекстуры в функции ядра - texture2d<float, access::write>, но они имеют разные форматы пикселей?

2. Является ли моя перезапись capturedImageFragmentShader в Отображение опыта AR с металлом каквычислительный шейдер так же прост, как я думал, или я что-то упустил?