iOS使用CoreML对视频进行实时图片分割

我使用的模型是DeepLabV3

可以直接对图片进行预测

    private func deeplabV3(img: CGImage) {
        guard let v3 = try? DeepLabV3(contentsOf: DeepLabV3.urlOfModelInThisBundle, configuration: MLModelConfiguration()) else { return }
        guard let input = try? DeepLabV3Input(imageWith: img) else { return }
        autoreleasepool {
            guard let output = try? v3.prediction(input: input) else { return }
            
            guard let size = delegate?.videoLayer?.bounds.size else { return }
            
            let width = size.width/513,
                height = size.height/513
            
            let path = CGMutablePath()
            for line in 0..<513 {
                for column in 0..<513 {
                    let green = output.semanticPredictions[[NSNumber(integerLiteral: line), NSNumber(integerLiteral: column)]]
                    if green != 0 {
                        path.addRect(CGRect(x: size.width-CGFloat(line)*width, y: CGFloat(column)*height, width: width, height: height))
                    }
                }
            }
            
            delegate?.handlerDeeplabV3(path: path)
        }
    }

我需要获取到视频中的图像

    func prediction(image: UIImage) {
        if let img = image.cgImage {
            deeplabV3(img: img)
        }
    }

但是通过AVCaptureSession得到的是CMSampleBuffer，所以这里需要做一些转换

    private func convertCGImage(sampleBuffer: CMSampleBuffer) {
        if let buffer = CMSampleBufferGetImageBuffer(sampleBuffer) {
            let ciimage = CIImage(cvPixelBuffer: buffer),
                context = CIContext()
            guard let cgimage = context.createCGImage(ciimage, from: ciimage.extent) else { return }
            deeplabV3(img: cgimage)
        }
    }

CMSampleBuffer的输出源就不做赘述了。

在测试的过程中出现一些问题

我在使用DeepLabV3直接进行预测时，由于视频中的图像是不断改变的，我需要不间断的进行预测，这样就会频繁调用

    guard let output = try? v3.prediction(input: input) else { return }

我发现会有内存泄漏的情况导致内存一直增加，这个问题现在还没解决，欢迎大佬们一起交流。