Core ML Model の画像出力を動画にする

17698 ワード
AVAssetWriter CreateML coreML Pixelbuffer avfoundation avfoundation テキストリンク

Create ML で Video Style Transfer モデルが追加された（WWDC2020）ということもあり、出力画像を動画にしたい場面もあるのではないでしょうか。

AVFoundation フレームワークで可能です。

ぼくはカメラの出力を一旦 Pixel Buffer コピーの配列にして保存してから、順次処理しました。
VideoCaptureOutputのフレームを参照し続けると処理が止まるので、PixelBuffer をディープコピーしてから使用した方がいいと思ってます。

手順

1,AVAssetWriterとを準備

サイズを Core ML の出力に合わせておきました。

func assetWriterSetting(){
    guard let url = try? FileManager.default.url(for: .documentDirectory, in: .userDomainMask, appropriateFor: nil, create: true).appendingPathComponent(fileName + ".mp4") else { print("nil"); return}
    print(url)
    let videoSettings = [
      AVVideoWidthKey: 256,
      AVVideoHeightKey: 256,
      AVVideoCodecKey: AVVideoCodecType.h264
    ] as [String: Any]

    videoAssetInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoSettings)
    pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: videoAssetInput, sourcePixelBufferAttributes: [kCVPixelBufferPixelFormatTypeKey as String: Int(kCVPixelFormatType_32BGRA)])
    frameNumber = 0
    do {
        try assetWriter = AVAssetWriter(outputURL: url, fileType: .mp4)
       assetWriter.add(videoAssetInput)
       assetWriter.startWriting()
       assetWriter.startSession(atSourceTime: CMTime.zero)
       } catch {
          print("could not start video recording ", error)
      }
}

2,PixelBufferを追加していく

VNCoreMLRequest の Completion Handler 内で Pixel Buffer を AVAssetWriterInputPixelBufferAdaptor に追加していきます。
ぼくは Core ML Helpers をつかって Multi Array から CGImage に変換したので、 CGImage を　Pixel Buffer に変換しています。

上記 AssetWriter のセッティングで sourcePixelBufferAttributes を BGRA にしているため、 RGB から BGR 色空間への変換もおこなっています。 AssetWriter のセッティングで RGB を設定する方法を知っている方がいたら、教えてください。

Pixel Buffer Poolがnilで返ってくる時は、AssetWriterの設定ミス（書き込み先URLが既に存在している。sourcePixelBufferAttributesで取扱不可のものを設定している）が原因のことが多いです。

let result = coreMLRequest.results?.first as! VNCoreMLFeatureValueObservation
let multiArray = result.featureValue.multiArrayValue
guard let cgimage = multiArray?.cgImage(min: -1, max: 1, channel: nil)?.toBGR() else {print("drop"); return}
guard let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool else {
    fatalError("Failed to allocate the PixelBufferPool")
}
var pixelBufferOut: CVPixelBuffer? = nil
CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pixelBufferPool, &pixelBufferOut)

guard let pixelBuffer = pixelBufferOut else {
    fatalError("Failed to create the PixelBuffer")
}

CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: CVOptionFlags(0)))

let context = CGContext(
    data: CVPixelBufferGetBaseAddress(pixelBuffer),
    width: cgimage.width,
    height: cgimage.height,
    bitsPerComponent: cgimage.bitsPerComponent,
    bytesPerRow: cgimage.bytesPerRow,
    space: CGColorSpaceCreateDeviceRGB(),
    bitmapInfo: cgimage.bitmapInfo.rawValue)
context?.draw(cgimage, in: cgimage.frame)

CVPixelBufferUnlockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: CVOptionFlags(0)))
if videoAssetInput.isReadyForMoreMediaData {

//処理速度が速すぎると、 isReadyForMoreMediaData が追いつかずに PixelBufferを追加できないことがあります。 適宜調整してください。

frameTime = CMTimeMake(value: Int64(Double(frameCount * fps) * durationForEachImage), timescale: Int32(fps))
pixelBufferAdaptor.append(pixelBuffer, withPresentationTime: frameTime)
frameCount += 1
print(frameTime)

CGImage のサイズ取得とBGR変換のextension

extension CGImage {
    var frame: CGRect {
        return CGRect(x: 0, y: 0, width: self.width, height: self.height)
    }

    func toBGR()->CGImage{
        let ciImage = CIImage(cgImage: self)
        let ctx = CIContext(options: nil)
        let swapKernel = CIColorKernel( source:
                                            "kernel vec4 swapRedAndGreenAmount(__sample s) {" +
                                            "return s.bgra;" +
                                            "}"
        )
        let ciOutput = swapKernel?.apply(extent: (ciImage.extent), arguments: [ciImage as Any])
        let cgOut:CGImage = ctx.createCGImage(ciOutput!, from: ciOutput!.extent)!
        return cgOut
    }
}

3,書き込み

if mlRequestsEnded {
   videoAssetInput.markAsFinished()
   assetWriter.endSession(atSourceTime: frameTime)    
   assetWriter.finishWriting(completionHandler: { [self] in
   print("comp")
}