初心者がYOLOをCore MLに変換してリアルタイムでオブジェクトカウンティングに挑戦してみた

$ pip install ultralytics

$ yolo detect predict model=yolo11n.pt source='https://ultralytics.com/images/bus.jpg'

$ yolo solutions count show=True

$ yolo export model=yolo11n.pt format=coreml nms=True

import UIKit
import AVFoundation
import CoreML
import Vision

class ViewController: UIViewController, AVCaptureVideoDataOutputSampleBufferDelegate {
    @IBOutlet weak var cameraView: UIImageView! // カメラのプレビュー用のビュー
    @IBOutlet weak var footerView: UIView! // カウンティング結果表示用のビュー
    private var textView: UILabel! // カウンティング結果表示用のラベル

    var frameInterval = 3 // フレーム間隔
    var frameCounter = 0 // フレームカウンター

    var captureSession: AVCaptureSession! // カメラのセッション
    var ciContext: CIContext! // CoreImageのコンテキスト
    // カメラの向きを調整するクラス（今回はコメントアウト（後述））
//    var rotationCoordinator: AVCaptureDevice.RotationCoordinator! 

    var yoloModel: VNCoreMLModel! // YOLOのモデル

    var mainCounter : [String: Int] = [:] // カウンターの初期配列

    override func viewDidLoad() {
        super.viewDidLoad()
        
        // カウンティング結果表示用のラベルの初期化
        textView = UILabel()
        textView.numberOfLines = 0
        textView.translatesAutoresizingMaskIntoConstraints = false
        footerView.addSubview(textView)
        NSLayoutConstraint.activate([
            textView.topAnchor.constraint(equalTo: footerView.topAnchor, constant: 10),
            textView.bottomAnchor.constraint(equalTo: footerView.bottomAnchor, constant: -10),
            textView.leadingAnchor.constraint(equalTo: footerView.leadingAnchor),
            textView.trailingAnchor.constraint(equalTo: footerView.trailingAnchor)
        ])
        
        // カメラのセットアップ
        setupCamera()

        ciContext = CIContext()
    }
}

/*!
 * @brief カメラのセットアップを行う
 */
func setupCamera() {
    captureSession = AVCaptureSession()

    // カメラのデバイスを取得
    guard let videoCaptureDevice = AVCaptureDevice.default(for: .video) else { return }

    // rotationCoordinatorの初期化（後述）
//    rotationCoordinator = AVCaptureDevice.RotationCoordinator(device: videoCaptureDevice, previewLayer: self.cameraView.layer)

    // ビデオ入力を設定
    let videoInput: AVCaptureDeviceInput

    do {
        videoInput = try AVCaptureDeviceInput(device: videoCaptureDevice)
    } catch {
      return
    }

    if (captureSession.canAddInput(videoInput)) {
        captureSession.addInput(videoInput)
    } else {
        return
    }

    // ビデオ出力を設定
    let videoOutput = AVCaptureVideoDataOutput()
    videoOutput.setSampleBufferDelegate(self, queue: DispatchQueue(label: "videoQueue"))
    if (captureSession.canAddOutput(videoOutput)) {
        captureSession.addOutput(videoOutput)
    } else {
        return
    }
    
    // セッションの開始
    captureSession.startRunning()
}

do {
    // モデルの初期化
    yoloModel = try VNCoreMLModel(for: yolo11n().model)
} catch {
    print("Model initialization error")
    return
}

struct DetectedObject {
    let observation: VNRecognizedObjectObservation // 検出結果オブジェクト
    let label: String // ラベル
    let confidence: Float // 信頼度
    let box: CGRect // オブジェクトの境界
}

/* !
 * @brief 画像からオブジェクトを検出する
 * 
 * @param [in] pixelBuffer 入力画像のピクセルバッファ
 * 
 * @return [DetectedObject] 検出されたオブジェクトの配列
 */
func detectObject(pixelBuffer: CVPixelBuffer) -> [DetectedObject]? {
    // ①リクエストを作成する
    let request = VNCoreMLRequest(model: yoloModel) { (request, error) in
        guard request.results is [VNRecognizedObjectObservation] else { return }
    }

    // リクエストハンドラーの作成
    let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer)
    do {
        // ②リクエストを実行する
        try handler.perform([request])
        guard let results = request.results as? [VNRecognizedObjectObservation] else { return nil }

        // ③実行結果を処理する
        var detectedObjects: [DetectedObject] = []

        // 画像サイズの取得
        let size = CGSize(width: CVPixelBufferGetWidth(pixelBuffer), height: CVPixelBufferGetHeight(pixelBuffer))
        
        for observation in results {
            // 信頼度が高いものだけ出力する（画面ビカビカ抑止のため目視で調整）
            if observation.confidence > 0.4 {
                guard let label = observation.labels.first?.identifier else { return nil }

/*
 * VNDetectedObjectObservationで取得できるboundingBoxは座標が左下原点で(0.0, 1.0)に正規化されています。
 * ここでは左上を原点とし、画像座標に変換しています。
 * 
 * VNImageRectForNormalizedRect()：正規化座標→画像座標に変換するメソッド。
 * sizeは座標の変換先である入力画像のサイズを示しています。
 * 以下参考：https://github.com/john-rocky/Yolov8-RealTime-iOS/blob/5226255dd1ea70243a1d0565515d7e3f0597e983/Yolov8-RealTime-iOS/ViewController.swift#L93
 */
                let vnBox = CGRect(x: observation.boundingBox.minX,
                                        y: 1 - observation.boundingBox.maxY, 
                                        width: observation.boundingBox.width,
                                        height: observation.boundingBox.height)
                let imageBox = VNImageRectForNormalizedRect(vnBox, Int(size.width), Int(size.height))

                // オブジェクトを配列に追加
                let object = DetectedObject(
                    observation: observation,
                    label: label,
                    confidence: observation.confidence,
                    box: imageBox)
                detectedObjects.append(object)
                print(object.label, observation.confidence)
            }
        }
        
        return detectedObjects
    } catch {
        print("Detection error")
        return nil
    }
}

/*!
 * @brief 検出されたオブジェクトを集計する
 *
 * @param [in] objects 検出されたオブジェクトの配列
 * 
 * @return 集計結果のテキスト
 */
func countObject(objects: [DetectedObject]) -> String? {
    var text = ""
    if objects.count > 0 {
        var objectCounter: [String: Int] = [:] // カウンターを初期化
        for object in objects {
            if objectCounter[object.label] != nil {
                objectCounter[object.label]! += 1
            } else {
                objectCounter[object.label] = 1
            }
        }

        // カウンターの値が変わった時だけ更新後のテキストを返す
        if mainCounter != objectCounter {
            mainCounter = objectCounter
            // カウンターをソートしてテキストに整形
            let sorted_counter = objectCounter.sorted(by: { $0.1 > $1.1 })
            for (label, count) in sorted_counter {
                text += "\(label): \(count)    "
                // person: 1  といった具合に表示される
            }
            return text
        }
    }
    return nil
}

/*!
 * @brief 検出したオブジェクトのボックスを描画する
 * 
 * @param [in] pixelBuffer 入力画像のピクセルバッファ
 * @param [in] objects 検出されたオブジェクトの配列
 * 
 * @return UIImage ボックス描画後の画像
 */
func drawBox(pixelBuffer: CVPixelBuffer, objects: [DetectedObject]) -> UIImage? {
        // 画像を変換
        let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
        let cgImage = self.ciContext.createCGImage(ciImage, from: ciImage.extent)!

        let size = ciImage.extent.size

        // CGContextを設定
        guard let cgContext = CGContext(data: nil,
                                        width: Int(size.width),
                                        height: Int(size.height),
                                        bitsPerComponent: 8,
                                        bytesPerRow: 4 * Int(size.width),
                                        space: CGColorSpaceCreateDeviceRGB(),
                                        bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else { return nil }
        cgContext.draw(cgImage, in: CGRect(origin: .zero, size: size))

        for object in objects {
            // CGImageに合わせてY座標を反転
            let invertedBox = CGRect(x: object.box.minX,
                                     y: size.height - object.box.maxY,
                                     width: object.box.width,
                                     height: object.box.height)

            // 線の設定と描画
            cgContext.setStrokeColor(UIColor.red.cgColor)
            cgContext.setLineWidth(5)
            cgContext.stroke(invertedBox)


        }

        guard let image = cgContext.makeImage() else { return nil }

        // 端末の向きに合わせて画像を回転させる
        var orientation: UIImage.Orientation!
        switch UIDevice.current.orientation {
            case .portrait:
                orientation = .right
            case .landscapeRight:
                orientation = .down
            case .landscapeLeft:
                orientation = .up
            default:
                orientation = .up
        }
        return UIImage(cgImage: image, scale: 1.0, orientation: orientation)
    }

/*!
 * @brief カメラの映像をキャプチャする
 * @param [in] output キャプチャした映像
 * @param [in] sampleBuffer サンプルバッファ
 * @param [in] connection キャプチャの接続
 */
func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    // frameInterval(3フレーム)ごとに処理を行う
    frameCounter += 1
    if frameCounter == frameInterval {
        frameCounter = 0

/* 
 * カメラの向きを縦向きに調整
 * videoRotationAngleForHorizonLevelCaptureを設定することで、撮影した画像が重力に対して水平になります。
 * フレーム出力ごとに設定しているので、端末を回転しても追従します。　
 */
//        connection.videoRotationAngle = rotationCoordinator.videoRotationAngleForHorizonLevelCapture

        // サンプルバッファからピクセルバッファを取得
        guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { return }

        // 物体検出と画像の描画を実行
        var text: String!
        var image: UIImage!
        if let objects = detectObject(pixelBuffer: pixelBuffer) {
            text = countObject(objects: objects)
            image = drawBox(pixelBuffer: pixelBuffer, objects: objects)
        }

        // UIの更新はメインキューで行う
        DispatchQueue.main.async {
            // テキスト更新
            if text != nil {
                self.textView.text = text
            }

            // 画像更新
            if image != nil {
                self.cameraView.image = image
            }
        }
    }
}