Swiftでアイマスの画像認識やってみる

はじめに

XcodeのCreateMLが簡単に画像認識をアプリケーションを組み込めるようなので試しにやってみました。

https://developer.apple.com/jp/machine-learning/create-ml/

上のようにやってくれるらしい。

作成したもの

作った pic.twitter.com/FuWqPuZsdH

— 清水 幸佑 (@thimi0412) December 2, 2019

環境

• Xcode Version 11.2.1

画像データの準備

何を分類するのか

りんごとかオレンジとかの分類だとなんかつまらないので、今回はアイドルマスターシンデレラガールズ カードギャラリーさんから画像をいただいて、島村卯月、渋谷凛、本田未央を分類してみようと思います。

学習用のデータとテスト用のデータのフォーマットはこんな感じ。
各ディレクトリの名前がラベルになるのでtrain_dataとtest_dataの名前を間違えないように気をつけましょう。

train_data
├── mio_honda
├── rin_shibuya
└── uzuki_shimamura

test_data
├── mio_honda
├── rin_shibuya
└── uzuki_shimamura

こんな感じで入ってます。

モデルの作成

プロジェクトの作成

Xcode > Open Developer Tool > Create MLから起動

プロジェクトのフォルダを決定しテンプレートを選択します。今回は画像認識なのでImage Classifierを選択。

学習開始

学習用データをTraining Dataに、テスト用データをTesting Dataにデータを追加しましょう。
追加したらTrainボタンを押して完了です。

学習が完了したらOutputにある.mlmodelファイルをドラッグして別の場所に保存しておきましょう。

アプリケーションに組み込む’

XcodeからCreate new Xcode projectでSingle View Appでプロジェクトを作成します。私はswiftには疎いのでとりあえず書いたコードを乗っけておきます。
ViewController.swiftとカメラを使用するのでinfo.plistをいじりました。

先ほど作成した.mlmodelをプロジェクトに追加してください。

info.plist

ViewController.swift

import UIKit
import AVKit
import CoreML
import Vision

class ViewController: UIViewController, AVCaptureVideoDataOutputSampleBufferDelegate {

    @IBOutlet weak var cameraDisplay: UIImageView!
    @IBOutlet weak var resultLabel: UILabel!
    override func viewDidLoad() {
        super.viewDidLoad()
        setUpCamere()
    }

    func setUpCamere() {
        guard let device = AVCaptureDevice.default(for: .video) else { return }
        guard let input = try? AVCaptureDeviceInput(device: device) else { return }
        
        let session = AVCaptureSession()
        session.sessionPreset = .hd4K3840x2160
        
        let previewLayler = AVCaptureVideoPreviewLayer(session: session)
        previewLayler.frame = view.frame
        cameraDisplay.layer.addSublayer(previewLayler)
        
        let output = AVCaptureVideoDataOutput()
        output.setSampleBufferDelegate(self, queue: DispatchQueue(label: "CameraOutput"))
        session.addInput(input)
        session.addOutput(output)
        session.startRunning()
    }
    

    
    func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
        guard let sampleBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { return }
        scanImage(buffer: sampleBuffer)
    }
    
    func scanImage(buffer: CVPixelBuffer) {
        // .mlmodelを読み込ませる
        guard let model = try? VNCoreMLModel(for: IdolClassifier_1().model) else { return }
        let request = VNCoreMLRequest(model: model) { request, _ in
            guard let results = request.results as? [VNClassificationObservation] else { return }
            guard let mostConfidentResult = results.first else { return }
            
            DispatchQueue.main.async {
                if mostConfidentResult.confidence >= 0.9 {
                    let confidenceText = "\n \(Int(mostConfidentResult.confidence * 100))% confidence"
                    switch mostConfidentResult.identifier {
                    case "uzuki_shimamura":
                        self.resultLabel.text = "uzuki_shimamura \(confidenceText)"
                    case "rin_shibuya":
                        self.resultLabel.text = "rin_shibuya\(confidenceText)"
                    case "mio_honda":
                        self.resultLabel.text = "mio_honda\(confidenceText)"
                    
                    default:
                        return
                    }
                } else {
                    self.resultLabel.text = "I don't know"
                }
            }
        }
        let requestHandler = VNImageRequestHandler(cvPixelBuffer: buffer, options: [:])
        do {
            try requestHandler.perform([request])
        } catch {
            print(error)
        }
        
    }
}