Распознавание лиц с помощью CoreML и ARKit

Evgeny Vladimirovich

Published in

NOP::Nuances of Programming

5 min readSep 11, 2019

Создаём приложение с одним окном

Для начала нам понадобится создать iOS проект «single view app» (прим: автор работает в Xcode).

Теперь у нас есть проект. Мы обойдёмся без сторибордов, кнопок и переключателей. Приложение будет выполнятся программно — чистый код 🤗.

Удалите main.storyboard и запишите следующий код в файл AppDelegate.swift:

func application(_ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]?) -> Bool {
        // Переопределение точки для настройки после запуска приложения.
        window = UIWindow(frame: UIScreen.main.bounds)
        window?.makeKeyAndVisible()
        let controller = ViewController()
        window?.rootViewController = controller
        
        return true
    }

Обязательно удалите информацию о сториборде Main из файлов.

Создаём сцену и добавляем её в Subview

У нас есть только один ViewController, который будет основной точкой входа для приложения.

На этом этапе нам нужно импортировать ARKit и создать экземпляр ARSCNView, который автоматически отображает видео в реальном времени с камеры устройства в качестве фона сцены. Кроме того, ARKit автоматически регулирует свою камеру SceneKit в соответствии с реальным движением устройства. Т.е. нам не нужен якорь для отслеживания положения объектов, которые мы добавляем к сцене.

Нам нужно задать границы экрана:

let sceneView = ARSCNView(frame: UIScreen.main.bounds)

В методе ViewDidLoad нам нужно настроить несколько вещей, например, делегат, а также вывести статистику кадров на экран:

self.view.addSubview(sceneView) // Добавление сцены в subview
sceneView.delegate = self // Установка делегата для view controller
sceneView.showsStatistics = true // Отображение статистики

Запуск сессии `ARFaceTrackingConfiguration`

Теперь нам нужно начать сессию ARFaceTrackingConfiguration. Эта конфигурация даёт нам доступ к фронтальной камере TrueDepth, которая есть только в iPhone X, Xs и Xr. Подробнее в документации Apple:

Конфигурация отслеживания лица определяет лицо пользователя в поле видимости фронтальной камеры устройства. При запуске этой конфигурации AR сессия обнаруживает лицо пользователя (если оно попадает в поле видимости фронтальной камеры) и добавляет в свой список якорей объект ARFaceAnchor, представляющий лицо. Каждый якорь передаёт информацию о положении, ориентации, топологии и других особенностях лица.
Источник: Apple

Метод ViewDidLoad выглядит так:

override func viewDidLoad() {
        super.viewDidLoad()
        self.view.addSubview(sceneView)
        sceneView.delegate = self
        sceneView.showsStatistics = true
        guard ARFaceTrackingConfiguration.isSupported else { return }
        let configuration = ARFaceTrackingConfiguration()
        configuration.isLightEstimationEnabled = true
        sceneView.session.run(configuration, options: [.resetTracking, .removeExistingAnchors])
    }

Обучаем модель

Есть несколько способов создать файл .mlmodel, который будет совместим с CoreML, вот основные:

Turicreate: это библиотека python, которая упрощает разработку моделей машинного обучения, и что ещё более важно, вы можете экспортировать свою модель в .mlmodel файл, который можно парсить с помощью Xcode.
MLImageClassifierBuilder(): здесь не требуется доп. библиотек, вам нужен только Xcode. Этот способ подойдёт для обучения относительно простых моделей.

Я создал несколько моделей, чтобы протестировать оба способа. У меня нет большого набора данных, поэтому я решил использовать MLImageClassifierBuilder(), 67 изображений своего лица и 260 изображений неизвестных лиц, которые я нашёл на unsplash.

Начнём с этого кода:

import CreateMLUI

let builder = MLImageClassifierBuilder()
builder.showInLiveView()

Я рекомендую установить максимальное количество итераций на 20 и добавить аугментацию изображений. Она создаст по четыре дополнительных экземпляра обрезанных изображений для каждого.

Захват фреймов с камеры и добавление их в модель

Нам нужно расширить наш ViewController с делегатом сцены — ARSCNViewDelegate. Нам понадобится два метода: один для обнаружения лица, а другой для обновления сцены при обнаружении лица.

Обнаружение лица:

func renderer(_ renderer: SCNSceneRenderer, nodeFor anchor: ARAnchor) -> SCNNode? {
        
        guard let device = sceneView.device else {
            return nil
        }
        
        let faceGeometry = ARSCNFaceGeometry(device: device)
        
        let node = SCNNode(geometry: faceGeometry)
        
        node.geometry?.firstMaterial?.fillMode = .lines
        
        return node
    }

К сожалению, сцена не обновляется, когда я открываю глаза или рот. Поэтому нам нужно каждый раз обновлять ее.

Обновление сцены:

func renderer(_ renderer: SCNSceneRenderer, didUpdate node: SCNNode, for anchor: ARAnchor) {
        
        guard let faceAnchor = anchor as? ARFaceAnchor,
            let faceGeometry = node.geometry as? ARSCNFaceGeometry else {
                return
        }
        
        faceGeometry.update(from: faceAnchor.geometry)
}

Захватываем всю геометрию лица и мэппинг. После этого обновляем узел.

Получаем фреймы с камеры:

Обратите внимание, что ARSCNView наследует от AVCaptureSession. Мы можем взять cvPixelFuffer и отправить его в нашу модель.

Вот простой способ сделать это с помощью атрибута sceneView:

guard let pixelBuffer = self.sceneView.session.currentFrame?.capturedImage else { return }

Добавляем фреймы в модель:

Теперь, когда мы можем обнаружить лицо и захватить каждый кадр с камеры, можно отправить их в нашу модель:

guard let model = try? VNCoreMLModel(for: FaceRecognition3().model) else {
            fatalError("Unable to load model")
        }
        
        let coreMlRequest = VNCoreMLRequest(model: model) {[weak self] request, error in
            guard let results = request.results as? [VNClassificationObservation],
                let topResult = results.first
                else {
                    fatalError("Unexpected results")
            }

            DispatchQueue.main.async {[weak self] in
                print(topResult.identifier)
            }
        }
        
        guard let pixelBuffer = self.sceneView.session.currentFrame?.capturedImage else { return }
        
        
        let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
        DispatchQueue.global().async {
            do {
                try handler.perform([coreMlRequest])
            } catch {
                print(error)
            }
        }

Отображаем имя над распознанным лицом

Последнее и, вероятно, самое сложное — это проецирование 3D-текста над распознанным лицом. Наша конфигурация не такая продвинутая, как ARWorldTrackingConfiguration, которая даёт доступ ко многим методам и классам. Мы используем только фронтальную камеру, и этим наши возможности ограничены.

Тем не менее, мы всё-таки можем спроецировать 3D-текст на экран, хотя он не будет реагировать на движение лица и меняться.

let text = SCNText(string: "", extrusionDepth: 2)
let font = UIFont(name: "Avenir-Heavy", size: 18)
text.font = font
let material = SCNMaterial()
material.diffuse.contents = UIColor.black
text.materials = [material]
text.firstMaterial?.isDoubleSided = true
        
let textNode = SCNNode(geometry: faceGeometry)
textNode.position = SCNVector3(-0.1, -0.01, -0.5)
textNode.scale = SCNVector3(0.002, 0.002, 0.002)
textNode.geometry = text

Теперь у нас есть объект SCNText. Он должен обновляться в соответствии с обнаруженным лицом. Мы должны добавить его в rootNode:

let coreMlRequest = VNCoreMLRequest(model: model) {[weak self] request, error in
            guard let results = request.results as? [VNClassificationObservation],
                let topResult = results.first
                else {
                    fatalError("Unexpected results")
            }

            DispatchQueue.main.async {[weak self] in
                print(topResult.identifier)
                if topResult.identifier != "Unknown" {
                    text.string = topResult.identifier
                    self!.sceneView.scene.rootNode.addChildNode(textNode)
                    self!.sceneView.autoenablesDefaultLighting = true
                }
            }
        }