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
使用真機測試，第一次保存沒有相冊權限會失敗。