Как отобразить маску поверх слоя камеры с помощью VNCoreMLFeatureValueObservation, возвращенного из моей .mlmodel? - PullRequest
Новая
       65

Как отобразить маску поверх слоя камеры с помощью VNCoreMLFeatureValueObservation, возвращенного из моей .mlmodel?

0 голосов
/ 09 июля 2020

Я просто обучил .mlmodel с semanti c сегментацией спецификациями задачи c. Я добавил AVCapturePreviewLayer с его делегатом.

В моем делегате я ловлю буфер образца и передаю его в Vision Framework как запрос изображения: 

func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    if let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) {
        let model = try! VNCoreMLModel(for: espnetv2_fp16().model)
        let request = VNCoreMLRequest(model: model) { request, error in
            DispatchQueue.main.async {
                if let results = request.results {
                    self.processResults(results)
                }
            }
        }
        request.imageCropAndScaleOption = .scaleFill
        let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
        try! handler.perform([request])
    }
}

func processResults(_ results: [Any]){
    segmentationOverlay.sublayers = nil // remove previos layers
    if let multiArray = (results as? [VNCoreMLFeatureValueObservation])?.first?.featureValue.multiArrayValue {
        let uiimage: UIImage! = nil
        // How do I turn multiArray into UIImage?
        let caLayer = CALayer()
        caLayer.frame = previewLayer.frame
        caLayer.contents = uiimage.cgImage
        caLayer.opacity = 0.5
        
        segmentationOverlay.addSublayer(caLayer)
        CATransaction.commit()
    }
}

Вопрос: Как превратить мой multiArray в правильный UIImage?

Важная информация:

  • Мой .mlnodel обучен распознавать гвозди в руке. На камере мне нужно, чтобы части изображения распознавались как гвозди и покрывали их любым цветом.

Результат должен быть примерно таким:

введите описание изображения здесь

...