CoreML物体识别 让AVCaptureVideoDataOutput和AVCaptureMovieFileOutput同时工作

如何让 AVCaptureVideoDataOutput 和 AVCaptureMovieFileOutput 同时工作

AVCaptureVideoDataOutput 是为了调用 CoreML 识别物体的数据流。我们通过 VNCoreMLRequest 来获取

guard let modelURL = Bundle.main.url(forResource: "YOLOv3FP16", withExtension: "mlmodelc") else {
    return NSError(domain: "VisionObjectRecognitionViewController", code: -1, userInfo: [NSLocalizedDescriptionKey: "Model file is missing"])
}
do {
    let visionModel = try VNCoreMLModel(for: MLModel(contentsOf: modelURL))

    let objectRecognition = VNCoreMLRequest(model: visionModel, completionHandler: { (request, error) in
        DispatchQueue.main.async(execute: {
            // perform all the UI updates on the main queue
            if let results = request.results {
                self.drawVisionRequestResults(results)
            }
        })
    })
    objectRecognition.imageCropAndScaleOption = .scaleFill
    self.requests = [objectRecognition]
} catch let error as NSError {
    print("Model loading went wrong: \(error)")
}

CoreML 物体识别过程

使用 CoreML 让手机直接识别物体需要如下几个步骤

1. 创建一个 AVCaptureSession 来获取摄像头的输入。

var deviceInput: AVCaptureDeviceInput!
session.addInput(deviceInput)

2. 创建一个 AVCaptureVideoDataOutput 来捕获输出。

session.addOutput(videoDataOutput)

3. 使用 Vision 框架里的 VNCoreMLModel 来获取 .mlmodel 模型对象。

let visionModel = try VNCoreMLModel(for: MLModel(contentsOf: modelURL))

4. 创建一个 VNRequest 来捕获输出图像用来物体识别。

let objectRecognition = VNCoreMLRequest(model: visionModel, completionHandler: { (request, error) in
    DispatchQueue.main.async(execute: {
        // perform all the UI updates on the main queue
        if let results = request.results {
            self.drawVisionRequestResults(results)
        }
    })
})
objectRecognition.imageCropAndScaleOption = .scaleFill
self.requests = [objectRecognition]

5. 在输出源的监听中加入 VNImageRequestHandler 来捕获每一帧。

let exifOrientation = exifOrientationFromDeviceOrientation()
let imageRequestHandler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, orientation: exifOrientation, options: [:])
do {
    try imageRequestHandler.perform(self.requests)
} catch {
    print(error)
}

让识别和录制同时工作

让 AVCaptureVideoDataOutput 和 AVCaptureMovieFileOutput 同时工作是为了实现边识别边录制的需求。我尝试了几种方法。

1. 把两个 output 放入 session

session.addOutput(movieFileOutput)
session.addOutput(videoDataOutput)

这样发现 session 无法正常工作。交换位置也不行，只能使用一种输出。
2. 定义两个 session 把 摄像头输入源赋给两个输入。

var deviceInput: AVCaptureDeviceInput!
recordSession.addInput(deviceInput)

因为识别和录制都需要摄像头，所以两个 session 都需要这个输入。一个 session 可以工作，另一个 session 同样无法正常工作。
测试后发现官方应该还不支持同时使用这两个输出。

其他尝试方法

1. AVCaptureVideoDataOutput+录制屏幕（会有权限弹窗，不好）

RPScreenRecorder.shared().startRecording { (err) in
    if let error = err {
        print(error)
    }

    DispatchQueue.main.asyncAfter(deadline: .now() + 5, execute:{
         ///延迟执行的代码

        RPScreenRecorder.shared().stopRecording { (previewCon, error) in
            if let errors = error {
                print(errors)
            }

            DispatchQueue.main.async(execute: {
                let url = previewCon!.movieURL
                print("has")
                print(url as Any)

            })
        }
    })
}

2. AVCaptureMovieFileOutput（无法获取每帧数据）
3. AVCaptureVideoDataOutput+图片收集转成mp4（视频过大时吃内存）

override func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
        return
    }

    if !CMSampleBufferDataIsReady(sampleBuffer) {
        return
    }
}

这个方法实际上是可行的，在拿到 CMSampleBuffer 后转化为图片，如果录制一个 10s 的短视频。如果需要长时间录制，内存会放不下，但也有改进空间，比如每10s生成一个短视频后清空图片缓存，最后生成多个短视频合起来。
但如果有直接写成视频的方法会更方便。所以使用下一个方案。
4. AVCaptureVideoDataOutput+图片异步处理成视频写入沙盒。

AVCaptureVideoDataOutput+图片异步处理成视频写入沙盒。

使用 AVAssetWriter 来写入视频。我找到了一个开源的视频写入类 PBJMediaWriter 作了少许修改。

1. 将一个视频写入地址传给输入来创建，可以在沙盒中创建一个指定文件夹来保存视频。

NSError *error = nil;
_assetWriter = [AVAssetWriter assetWriterWithURL:outputURL fileType:(NSString *)kUTTypeMPEG4 error:&error];
if (error) {
    DLog(@"error setting up the asset writer (%@)", error);
    _assetWriter = nil;
    return nil;
}

2. 创建两个 AVAssetWriterInput 来记录视频和音频。

AVAssetWriterInput *_assetWriterAudioInput;
AVAssetWriterInput *_assetWriterVideoInput;

3. 在获取视频帧后检查初始化，并写入每一帧。

override func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
        return
    }

    if !CMSampleBufferDataIsReady(sampleBuffer) {
        return
    }

    if !writer.isVideoReady {
        writer.setupMediaWriterVideoInput(with: sampleBuffer)
        return
    }
    if isRecording {
        if writer.isVideoReady {
            writer.write(sampleBuffer, withMediaTypeVideo: true)
        }
    }
}

4. 最后在你完成时，比如按了停止按钮后生成视频。

self.writer.finishWriting {
    print(self.writer.error as Any)
    self.saveVideoToAlbum(videoUrl: self.writer.outputURL)
}

这期间的录制都不会影响识别的过程，因为 pixelBuffer 可以共用，同时给物体识别和视频录制。
让 AVCaptureVideoDataOutput 和 AVCaptureMovieFileOutput 同时工作我们没有实现，但可以通过这个方案实现一样的效果。

demo：CoreMLRecord
使用真机测试，第一次保存没有相册权限会失败。