在 iOS 相机 App 训练 Create ML Style Transfer 模型!

522 阅读10分钟

样式转换 (Style Transfer) 是一个非常热门的深度学习课题,它可以让我们在一个图像上套用另一个图像的视觉样式,来改变图像的构图。

透过神经样式转换模型 (Neural Style Transfer Model),我们可以建构很多令人惊艳的东西,像是建构艺术照片编辑器、或是应用最新设计的样式为游戏设计赋予新样貌等。它非常方便,也可以使用资料扩充。

在 WWDC 2020,Create ML(Apple 的模型建构框架)加入了样式转换模型 (Style Transfer Model),功能大大提升。儘管 Xcode 12 随附这项更新,但你还需要 macOS Big Sur(在撰写本篇文章时还是 Beta 版本)来训练样式转换模型。

初探 Create ML 的样式转换

现在,Create ML 让我们可以使用 Macbook 直接训练样式转换模型。你可以训练图像与影像的样式转换卷积神经网路 (Style Transfer Convolutional Neural Networks),影像的转换只会使用有限的卷积过滤器,让它最佳化实时的图像处理。

在开始之前,你需要三项东西:

  • 一个样式图像(也可以叫做样式参考图像)。通常来说,你可用著名画作或是抽象艺术图像,来让你的模型学习风格。在范例中,我们会在第一个模型裡使用铅笔素描图像(参看下面的萤幕截图)。
  • 一张验证图像,它为我们在训练过程中可视化模型品质。
  • 内容图像的资料集,这是我们的训练资料。为了取得最佳结果,这裡使用的图像目录,最好与你执行推算时使用的图像相似。

在本篇文章中,我将使用这个名人图集作为内容图像。

在训练模型之前,让我们先来看看 Create ML 样式转换设定页籤。

style-transfer-1

以下的验证图像显示每个迭代 (iteration) 区间内,即时套用的样式转换:

style-transfer-2

让我们来看看样式损失 (Style Loss) 与内容损失(Content Loss) 的折线图,这是了解样式和内容图像之间平衡的指标。通常来说,样式损失应该随时间减少,这表示神经网路正在学习调用样式图像的艺术特徵。

虽然预设模型参数运作得不错,不过 Create ML 允许我们因应特定条件来作客制化。

低样式强度参数仅使用样式图像微调部分背景,保持图像主体完整。如果把样式强度参数设定为高,就会在图像边缘上给予更多样式纹理。

同样地,粗糙的样式密度会使用样式图像的高阶细节(训练这类的模型速度要快得多),而细密的密度则让模型学习细微的细节。

Create ML 样式转换模型训练预设迭代为 500 次,这对于大部分使用案例来说都是理想的数值。迭代是完成一个阶段所需的批次数量,一个阶段等于整个资料集的一个训练週期。例如,如果训练的资料集是由 500 张图像组成,而批次的数量是 50,那就表示 10 次迭代完成一个阶段(注意:Create ML 模型训练并不会告诉你批次大小。)

今年,Create ML 还引入了另外一个新功能:模型快照 (Model Snapshots)。这项功能让我们可以在模型训练期间撷取 Core ML 模型,并汇出到 App 裡。然而,来自快照的模型并未最佳化档案大小,而且明显地大于完成训练后的模型。具体来说,我撷取的 Core ML 模型快照大小为 5 – 6 MB,而最终模型大小为 596 KB。

我比较了不同迭代中的模型快照,结果如下:

style-transfer-3

请注意,在其中一张图像中,样式无法构成完整的图像。这是因为使用的样式图像尺寸较小,所以网路无法学习足够的样式资讯,导致合成图像的品质不佳。

理想来说,样式图像应为 512 px 以上,这样才能确保较好的结果。

目标

在接下来的部分,我们将建构一个 iOS App 来实时运行样式转换模型。让我们先看看实作步骤的概览:

  • 从三个影像样式转换神经网路模型来分析结果。其中一个是用预设参数来训练的,其他则用高与低的样式强度参数训练。
  • 在 iOS App 裡使用 AVFoundation 来实作客制化相机。
  • 在实时相机上执行建构好的 Core ML 模型。我们会使用 Vision 请求来快速在萤幕上执行、推算、及绘制样式相机影格。
  • 检视 CPU、GPU 与神经引擎 (Neural Engine) 的结果。

要找一幅具有良好艺术效果的样式图像非常棘手。幸运的是,我简单地在 Google 搜寻找到这张图像

分析不同强度的样式转换模型

我已经用同一组资料集训练好了三个模型。来看看结果:

Image for post

如你所见,使用低强度模型,样式图像几乎没有对内容图像造成影响;但使用高强度模型,内容图像的边缘就有更多样式效果。

现在,我们已经准备好将模型(大约 0.5 MB 的大小)转移到 App 中。

Create ML 会让我们预览影像结果,不过速度非常地慢。幸运的是,我们很快就会在范例 App 裡实时显示它们。

AVFoundation 基础

AVFoundation 是 Apple 一个可高度客制化的多媒体内容框架。你可以绘画客制的遮罩、微调相机设定、使用深度输出进行照片切割、以及分析影格。

我们主要专注于分析影格、对影格进行样式转换、并在 UIImageView 上显示它们,来建构一个实时相机功能。(你也可以使用 Metal 进行最佳化,但为简单起见,本次教学中将跳过这个部分。)

基本来说,要建构一个客制化相机,我们需要使用以下元件:

  • AVCaptureSession:这个元件用来管理整个相机的 Session。它的功能包括了存取 iOS 装置输入、及传送资料到输出装置。AVCaptureSession 也让我们为不同的 Capture Session 定义 [Preset](https://developer.apple.com/documentation/avfoundation/avcapturesession/preset) 类型。
  • AVCaptureDevice:这个元件让我们选择前置或后置镜头。我们可以选择预设设定,或是使用 AVCaptureDevice.DiscoverySession 来过滤及选择特定的硬体功能,像是原深感测 (TrueDepth) 或广角 (WideAngle) 相机。
  • AVCaptureDeviceInput:这个元件提供从 Capture Device 撷取的多媒体来源,并将其送往 Capture Session。
  • AVCaptureOutput:这是一个抽象类别,提供输出内容到 Capture Session。它也让我们处理相机方向。我们可以设定多个输出(像是相机与麦克风)。比如说,如果你要拍摄相片及电影,那麽就要添加 AVCaptureMovieFileOutputAVCapturePhotoOutput。在我们的范例中,我们将使用 AVCaptureVideoDataOutput,来提供影像影格供我们处理。
  • AVCaptureVideoDataOutputSampleBufferDelegate:这是一个协定,让我们可以在 didOutput 协定方法裡存取每个影格缓衝 (Frame Buffer)。要启动接受影格,我们需要调用 AVCaptureVideoDataOutput 裡的 setSampleBufferDelegate
  • AVCaptureVideoPreviewLayer:这基本上是一个 CALayer,它视觉地显示来自 Capture Session 输出的实际相机画面。我们可以使用遮罩及动画来转变图层。对于 Sample Buffer 委託方法的运作来就,这个设定十分重要。

设定我们的客制化相机

首先,把 NSCameraUsageDescription 相机权限添加到 Xcode 专案的 info.plist

现在,是时候在 ViewController.swift 建立一个 AVCaptureSession

let captureSession = AVCaptureSession()
captureSession.sessionPreset = AVCaptureSession.Preset.medium

接著,在 AVCaptureDevice 实例裡可使用的相机类型列表中,让我们过滤并选择广角相机。然后,将它添加至 AVCaptureInput,并将其设定在 AVCaptureSession 上:

let availableDevices = AVCaptureDevice.DiscoverySession(deviceTypes: [.builtInWideAngleCamera], mediaType: AVMediaType.video, position: .back).devices

let availableDevices = AVCaptureDevice.DiscoverySession(deviceTypes: [.builtInWideAngleCamera], mediaType: AVMediaType.video, position: .back).devices
do {
    if let captureDevice = availableDevices.first {
        captureSession.addInput(try AVCaptureDeviceInput(device: captureDevice))
    }
} catch {
    print(error.localizedDescription)
}

现在,输入已经设定好了,让我们添加影像输出到 captureSession 上:

let videoOutput = AVCaptureVideoDataOutput()

videoOutput.alwaysDiscardsLateVideoFrames = true

videoOutput.setSampleBufferDelegate(self, queue: DispatchQueue(label: "videoQueue"))

if captureSession.canAddOutput(videoOutput){
 captureSession.addOutput(videoOutput)
}

alwaysDiscardsLateVideoFrames 属性可以确保延迟的影格会被丢弃,以减少延迟的情况。

最后,加入以下程式码来避免相机转向:

guard let connection = videoOutput.connection(with: .video) 
else { return }

guard connection.isVideoOrientationSupported else { return }
connection.videoOrientation = .portrait

备注:为了确保所有的转向,你需要基于装置现在的方向设定 videoOrientation。在本教学的最后会有程式码范例。

现在,我们终于可以添加预览图层,并启动 Camera Session:

let previewLayer = AVCaptureVideoPreviewLayer(session: captureSession)

view.layer.addSublayer(previewLayer)

captureSession.startRunning()

以下是我们刚建立的 configureSession() 方法:

Image for post

运用 Vision 为 Core ML 模型实时执行样式转换

现在来到了机器学习的部分。我们会使用 Vision 框架,为样式转换模型的输入图像做前处理。

让 ViewController 符合 AVCaptureVideoDataOutputSampleBufferDelegate 协定后,我们可以使用以下的方法取得每个影格:

extension ViewController: AVCaptureVideoDataOutputSampleBufferDelegate{

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

}

从上面的 Sample Buffer 实例,我们可以取得 CVPixelBuffer 实例,并且将它传送到 Vision 请求:

guard let model = try? VNCoreMLModel(for: StyleBlue.init(configuration: config).model) else { return }
let request = VNCoreMLRequest(model: model) { (finishedRequest, error) in
    guard let results = finishedRequest.results as? [VNPixelBufferObservation] else { return }
    guard let observation = results.first else { return }
    DispatchQueue.main.async(execute: {
        self.imageView.image = UIImage(pixelBuffer: observation.pixelBuffer)
    })
}
guard let pixelBuffer: CVPixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { return }
try? VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:]).perform([request])

VNCoreMLModel 作为容器,我们以下列方式在当中实例化 Core ML 模型:

StyleBlue.init(configuration: config).model

configMLModelConfiguration 型别的实例,用来定义 computeUnits 属性。这个属性可让我们设定 cpuOnlycpuAndGpu 或是 all(Neural Engine),以在我们想要的装置硬体上执行。

let config = MLModelConfiguration()
switch currentModelConfig {
case 1:
config.computeUnits = .cpuOnly
case 2:
config.computeUnits = .cpuAndGPU
default:
config.computeUnits = .all
}

备注:我们已设定了一个 UISegmentedControl UI 控件,让我们可以切换上面的模型配置。

我们在 VNCoreMLRequest 中传递 VNCoreMLModelVNCoreMLRequest 会回传 VNPixelBufferObservation 型别的观察 (Observations)。

VNPixelBufferObservationVNObservation 的子类别,它会回传 CVPixelBuffer 的影像输出。

我们利用下面的 Extension 来转换 CVPixelBuffer 为 UIImage,并将它绘制在萤幕上。

extension UIImage {
    public convenience init?(pixelBuffer: CVPixelBuffer) {
        var cgImage: CGImage?
        VTCreateCGImageFromCVPixelBuffer(pixelBuffer, options: nil, imageOut: &cgImage)
        if let cgImage = cgImage {
            self.init(cgImage: cgImage)
        } else {
            return nil
        }
    }
}

呼~!我们已经建立好实时样式转换 iOS App 了。

以下是在 iPhone SE 上执行 App 的成果:

image

image

image

从上面的 gif 中可见,顶部的分类选择器让你在两个模型之间挑选。StarryBlue 是以预设样式强度训练,而 Strong 页籤则是执行第二个高样式强度的模型。

让我们看看,在 Neural Engine 上运行时,样式转换预测是如何实时进行的。

备注:因为 gif 档案大小以及画质限制,我*录制了一段影片**,来示范 CPU、GPU、以及 Neural Engin 上的实时样式转换。影片会比上面的 gif 更加顺畅。*

你可以在 这个 GitHub Repository 中,参考这个 Core ML 样式转换模型 App 的完整程式码。

iOS 14 的 Core ML 引入了模型加密 (Model Encryption),因此理论上可以我加密这个模型。但是本著学习精神,我会免费提供上述创建的模型。

总结

机器学习的未来显然是不需要程式码的, MakeMLFritz AI、以及 Apple 的 Create ML 等平台将引领趋势,并提供易于使用的工具及平台,以快速训练可在手机上使用的机器学习模型。

Create ML 今年还引入支援人类活动分类的模型训练,但我相信样式转换会是快速地被 iOS 开发者所用应用的功能。如果你希望建构一个包含多个样式图像的单一模型,可以使用 Turi Create

现在,你可以使用神经样式转换神经网路,以零成本(除非你要把 App 上传到 App Store)建构基于 AI、类似 Prisma 的 App。

样式转换也适用于 ARKit 物件,给予它们一个完全不同的外观。我们将在下篇教学讨论这个部分,请密切留意我们的更新。

本篇文章到此为止,感谢你的阅读。

文末推荐:iOS热门文集