【Swift】簡単なリアルタイムに画像を分類するアプリを作ってみた

import AVFoundation

protocol VideoCaptureDelegate: AnyObject {
    func didSet(_ previewLayer: AVCaptureVideoPreviewLayer)
    func didCaptureFrame(from imageBuffer: CVImageBuffer)
}

class VideoCapture: NSObject {

    weak var delegate: VideoCaptureDelegate?

    // AVの出入力のキャプチャを管理するセッションオブジェクト
    private let captureSession = AVCaptureSession()

    // ビデオを記録し、処理を行うためにビデオフレームへのアクセスを提供するoutput
    private let videoOutput = AVCaptureVideoDataOutput()

    // カメラセットアップとフレームキャプチャを処理する為のDispathQueue
    private let sessionQueue = DispatchQueue(label: "object-detection-queue")

    func startCapturing() {

        // capture deviceのメディアタイプを決め、
        // 何からインプットするかを決める
        guard let captureDevice = AVCaptureDevice.default(for: .video),
              let deviceInput = try? AVCaptureDeviceInput(device: captureDevice)
        else { return }

        // captureSessionにdeviceInputを入力値として入れる
        captureSession.addInput(deviceInput)

        // キャプチャセッションの開始
        captureSession.startRunning()

        // ビデオフレームの更新ごとに呼ばれるデリゲートをセット
        videoOutput.setSampleBufferDelegate(self, queue: sessionQueue)

        // captureSessionから出力を取得するためにdataOutputをセット
        captureSession.addOutput(videoOutput)

        // captureSessionをUIに描画するためにPreviewLayerにsessionを追加
        let previewLayer = AVCaptureVideoPreviewLayer(session: captureSession)
        delegate?.didSet(previewLayer)
    }

    func stopCapturing() {
        // キャプチャセッションの終了
        captureSession.stopRunning()
    }
}

weak var delegate: VideoCaptureDelegate?

// AVの出入力のキャプチャを管理するセッションオブジェクト
private let captureSession = AVCaptureSession()

// ビデオを記録し、処理を行うためにビデオフレームへのアクセスを提供するoutput
private let videoOutput = AVCaptureVideoDataOutput()

// カメラセットアップとフレームキャプチャを処理する為のDispathQueue
private let sessionQueue = DispatchQueue(label: "object-detection-queue")

func startCapturing() {

    // capture deviceのメディアタイプを決め、
    // 何からインプットするかを決める
    guard let captureDevice = AVCaptureDevice.default(for: .video),
          let deviceInput = try? AVCaptureDeviceInput(device: captureDevice)
    else { return }

    // captureSessionにdeviceInputを入力値として入れる
    captureSession.addInput(deviceInput)

    // キャプチャセッションの開始
    captureSession.startRunning()

    // ビデオフレームの更新ごとに呼ばれるデリゲートをセット
    videoOutput.setSampleBufferDelegate(self, queue: sessionQueue)

    // captureSessionから出力を取得するためにdataOutputをセット
    captureSession.addOutput(videoOutput)

    // captureSessionをUIに描画するためにPreviewLayerにsessionを追加
    let previewLayer = AVCaptureVideoPreviewLayer(session: captureSession)
    delegate?.didSet(previewLayer)
}

let previewLayer = AVCaptureVideoPreviewLayer(session: captureSession)

func stopCapturing() {
    // キャプチャセッションの終了
    captureSession.stopRunning()
}

extension VideoCapture: AVCaptureVideoDataOutputSampleBufferDelegate {

    func captureOutput(_ output: AVCaptureOutput,
                       didOutput sampleBuffer: CMSampleBuffer,
                       from connection: AVCaptureConnection) {

        // フレームからImageBufferに変換
        guard let imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer)
        else { return }

        delegate?.didCaptureFrame(from: imageBuffer)
    }
}

import CoreML
import Vision

protocol Resnet50ModelManagerDelegate: AnyObject {
    func didRecieve(_ observation: VNClassificationObservation)
}

class Resnet50ModelManager: NSObject {

    weak var delegate: Resnet50ModelManagerDelegate?

    func performRequet(with imageBuffer: CVImageBuffer) {
        // 機械学習モデル
        guard let model = try? VNCoreMLModel(for: Resnet50(configuration: .init()).model)
        else { return }

        // フレーム内で機械学習モデルを使用した画像分析リクエスト
        let request = VNCoreMLRequest(model: model) { request, error in
            if let error = error {
                print(error.localizedDescription)
                return
            }

            guard let results = request.results as? [VNClassificationObservation],
                  let firstObservation = results.first
            else { return }

            self.delegate?.didRecieve(firstObservation)
        }

        // imageRequestHanderにimageBufferをセット
        let imageRequestHandler = VNImageRequestHandler(cvPixelBuffer: imageBuffer)
        // imageRequestHandlerにrequestをセットし、実行
        try? imageRequestHandler.perform([request])
    }
}

weak var delegate: Resnet50ModelManagerDelegate?

guard let model = try? VNCoreMLModel(for: Resnet50(configuration: .init()).model)
else { return }

// フレーム内で機械学習モデルを使用した画像分析リクエスト
let request = VNCoreMLRequest(model: model) { request, error in
    if let error = error {
        print(error.localizedDescription)
        return
    }

guard let results = request.results as? [VNClassificationObservation],
      let firstObservation = results.first
else { return }

self.delegate?.didRecieve(firstObservation)
}

// imageRequestHanderにimageBufferをセット
let imageRequestHandler = VNImageRequestHandler(cvPixelBuffer: imageBuffer)
// imageRequestHandlerにrequestをセットし、実行
try? imageRequestHandler.perform([request])

typealias CVPixelBuffer = CVImageBuffer

import UIKit
import AVFoundation
import Vision

class RealTimeImageClassficationViewController: UIViewController {

    @IBOutlet private weak var previewView: UIView!
    @IBOutlet private weak var observationLabel: UILabel!

    private let videoCapture = VideoCapture()
    private let resnet50ModelManager = Resnet50ModelManager()

    override func viewDidLoad() {
        super.viewDidLoad()
        resnet50ModelManager.delegate = self
        videoCapture.delegate = self
        videoCapture.startCapturing()
    }

    override func viewWillDisappear(_ animated: Bool) {
        super.viewWillDisappear(animated)
        videoCapture.stopCapturing()
    }
}

extension RealTimeImageClassficationViewController: VideoCaptureDelegate {

    // previewLayerがセットされた時に呼ばれる
    func didSet(_ previewLayer: AVCaptureVideoPreviewLayer) {
        previewView.layer.addSublayer(previewLayer)
        previewLayer.frame = previewView.frame
    }

    // フレームがキャプチャされる度に呼ばれる
    func didCaptureFrame(from imageBuffer: CVImageBuffer) {
        resnet50ModelManager.performRequet(with: imageBuffer)
    }
}

extension RealTimeImageClassficationViewController: Resnet50ModelManagerDelegate {

    // 画像分析リクエストから観測値を受け取る度に呼ばれる
    func didRecieve(_ observation: VNClassificationObservation) {

        DispatchQueue.main.async {
            self.observationLabel.text = "\(observation.confidence.convertPercent)%の確率で、\(observation.identifier)"
            print(observation.identifier, observation.confidence)
        }
    }
}

override func viewDidLoad() {
    super.viewDidLoad()
    resnet50ModelManager.delegate = self
    videoCapture.delegate = self
    videoCapture.startCapturing()
}

override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    videoCapture.stopCapturing()
}

func didSet(_ previewLayer: AVCaptureVideoPreviewLayer) {
    previewView.layer.addSublayer(previewLayer)
    previewLayer.frame = previewView.frame
}

func didCaptureFrame(from imageBuffer: CVImageBuffer) {
    resnet50ModelManager.performRequet(with: imageBuffer)
}

func didRecieve(_ observation: VNClassificationObservation) {

    DispatchQueue.main.async {
        self.observationLabel.text = "\(observation.confidence.convertPercent)%の確率で、\(observation.identifier)"
        print(observation.identifier, observation.confidence)
    }
}

extension VNConfidence {
    var convertPercent: Int {
        return Int(self * 100)
    }
}