将 MLMultiArray 转换为图像或 OpenGL / 金属纹理

Question

I'm trying to do background segmentation of a live video using CoreML. I used DeepLabV3 as provided by Apple. The model works ok, even though it already takes 100ms to process a 513x513 image. I then want to display the output, which is a 513x513 array of int32. Converting it in an image as done in CoreMLHelpers takes 300ms and I'm looking for a much faster way to display the results. I was thinking that maybe it'd be faster to somehow dump this to a OpenGL or Metal texture.

What is the best way to handle MLMultiArray for live inputs?

Answer 1

My answer is based on processing the MLMultiArray in Metal

Create an MTLBuffer:

let device = MTLCreateSystemDefaultDevice()!
let segmentationMaskBuffer: MTLBuffer = self.device.makeBuffer(length: segmentationHeight * segmentationWidth * MemoryLayout<Int32>.stride)

Copy MLMultiArray to MTLBuffer:

memcpy(segmentationMaskBuffer.contents(), mlOutput.semanticPredictions.dataPointer, segmentationMaskBuffer.length)

Setup Metal related variables:

let commandQueue = device.makeCommandQueue()!
let library = device.makeDefaultLibrary()!
let function = library.makeFunction(name: "binaryMask")!
let computePipeline = try! device.makeComputePipelineState(function: function)

create a struct for segmentation size:

let segmentationWidth = 513
let segmentationHeight = 513

struct MixParams {
    var width: Int32 = Int32(segmentationWidth)
    var height: Int32 = Int32(segmentationHeight)
}

create a output texture:

let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .bgra8Unorm, width: width, height: height, mipmapped: false)
textureDescriptor.usage = [.shaderRead, .shaderWrite]
let outputTexture = device.makeTexture(descriptor: textureDescriptor)!

pass the mtlbuffer, outputtexture to the kernal function:

let buffer = commandQueue.makeCommandBuffer()!
let maskCommandEncoder = buffer.makeComputeCommandEncoder()!
maskCommandEncoder.setTexture(outputTexture, index: 1)
maskCommandEncoder.setBuffer(segmentationBuffer, offset: 0, index: 0)
maskCommandEncoder.setBytes(&params, length: MemoryLayout<MixParams>.size, index: 1)
let w = computePipeline.threadExecutionWidth
let h = computePipeline.maxTotalThreadsPerThreadgroup / w
let threadGroupSize = MTLSizeMake(w, h, 1)
let threadGroups = MTLSizeMake(
          (depthWidth  + threadGroupSize.width  - 1) / threadGroupSize.width,
          (depthHeight + threadGroupSize.height - 1) / threadGroupSize.height, 1)
maskCommandEncoder.setComputePipelineState(computePipeline)
maskCommandEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
maskCommandEncoder.endEncoding()

write your kernel function in Shaders.metal file:

#include <metal_stdlib>
using namespace metal;
#include <CoreImage/CoreImage.h>

struct MixParams {
    int segmentationWidth;
    int segmentationHeight;
};

static inline int get_class(float2 pos, int width, int height, device int* mask) {
    const int x = int(pos.x * width);
    const int y = int(pos.y * height);
    return mask[y*width + x];
}

static float get_person_probability(float2 pos, int width, int height, device int* mask) {   
    return get_class(pos, width, height, mask) == 15;
}

kernel void binaryMask(
                      texture2d<float, access::write> outputTexture [[texture(1)]],
                      device int* segmentationMask [[buffer(0)]],
                      constant MixParams& params [[buffer(1)]],
                      uint2 gid [[thread_position_in_grid]])
{
    float width = outputTexture.get_width();
    float height = outputTexture.get_height();
    
    if (gid.x >= width ||
        gid.y >= height) return;
    
    
    const float2 pos = float2(float(gid.x) / width,
                              float(gid.y) / height);
    
    const float is_person = get_person_probability(pos, params.segmentationWidth,
                                                   params.segmentationHeight,
                                                   segmentationMask);
    
    float4 outPixel;
    
    if (is_person < 0.5f) {
        outPixel = float4(0.0,0.0,0.0,0.0);
    } else {
        outPixel = float4(1.0,1.0,1.0,1.0);
    }
    
    outputTexture.write(outPixel, gid);
}

Finally get the ciimage from output texture after encoding is finished:

let kciOptions: [CIImageOption: Any] = [CIImageOption.colorSpace: CGColorSpaceCreateDeviceRGB()]
let maskIMage = CIImage(mtlTexture: outputTexture,options: kciOptions)!.oriented(.downMirrored)

Answer 2

Instead of outputting an MLMultiArray you can change the model to make it output an image of type CVPixelBuffer . Then you can use CVMetalTextureCacheCreateTextureFromImage to turn the pixel buffer into an MTLTexture . (I think this works but I don't recall if I ever tried it. Not all pixel buffer objects can be turned into textures and I'm not sure if Core ML outputs a CVPixelBuffer object with the "Metal compatibility flag" turned on.)

Alternatively, you can write a compute kernel that takes in the MLMultiArray and converts it to a texture, which then gets drawn into a Metal view. This has the advantage that you apply all kinds of effects to the segmentation map in the compute kernel at the same time.