WWDC Quick Look 💓 By SwiftGGTeam
Explore the machine learning development experience

Explore the machine learning development experience

观看原视频

Highlight

Session 以黑白老照片自动上色为例,串起模型搜索、PyTorch 到 Core ML 转换、Xcode 性能评估、Swift 集成和 Instruments 调优,并把 90ms 的静态图片模型优化成约 16ms 的实时摄像头方案。


核心内容

很多 App 想加入机器学习功能,第一步并不是写模型代码。开发者先要判断:这个问题是否适合机器学习,模型从哪里来,怎样放进 Apple 平台,转换后是否正确,目标设备上是否够快。

这场 Session 用一个具体任务回答这些问题:把地下室旧盒子里的黑白家庭照片自动上色。人工修图需要专业摄影师和图像编辑工具,机器学习可以把流程缩短到几秒。

Apple 的示例流程很直接。先找到一个开源 PyTorch 模型 Colorizer;再用 coremltools 转成 Core ML;接着用 Python 验证转换结果;然后用 Xcode 14 的 Core ML Performance Report 看设备预测时间;最后用 Core Image 和 Core ML 把模型接进 Swift App。

静态图片版本已经够用。Performance Report 显示,原模型在 iPad Pro M1 和 iPadOS 16 上预测约 90ms。需求变成实时摄像头后,这个数字不够了:30fps 摄像头约每 30ms 产生一帧,90ms 会让每次上色期间丢掉两到三帧。

后半段因此进入优化流程。演讲者改动模型架构,用 PyTorch on Metal 在 Apple Silicon 上重新训练,再转换、验证、集成。新模型完全跑在 Neural Engine 上,预测时间约 16ms。最后用 Instruments 的 Core ML 模板发现 App 排队提交预测和误设 60fps 的问题,改成上一帧预测结束后再处理下一帧,并把摄像头设回 30fps。

详细内容

1. 把 RGB 图像拆成模型需要的 L 通道(03:06

Colorizer 使用 Lab 色彩空间。Lab 有三个通道:L 表示亮度,a 和 b 表示颜色分量。输入黑白图时,模型只需要 L 通道,然后预测新的 a、b 通道。

from skimage import color

in_lab = color.rgb2lab(in_rgb)
in_l = in_lab[:,:,0]

关键点:

  • color.rgb2lab(in_rgb) 把 RGB 图像转换到 Lab 色彩空间。
  • in_lab[:,:,0] 取出第一个通道,也就是亮度 L。
  • a、b 两个颜色通道在输入阶段被丢弃,留给模型预测。

模型输出后,要把输入 L 和预测出的 a、b 重新拼回 Lab,再转回 RGB。

from skimage import color
import numpy as np
import torch

out_lab = torch.cat((in_l, out_ab), dim=1)
out_rgb = color.lab2rgb(out_lab.data.numpy()[0,…].transpose((1,2,0)))

关键点:

  • torch.cat((in_l, out_ab), dim=1) 把原始亮度和模型输出的颜色分量按通道拼接。
  • out_lab.data.numpy() 把张量转成 NumPy 数据,供图像处理库使用。
  • transpose((1,2,0)) 把通道维移到最后,匹配 lab2rgb 期望的图像布局。
  • color.lab2rgb(...) 得到最终彩色图像。

2. 用 coremltools 转换 PyTorch 模型(03:56

找到开源模型后,下一步是让它在 Apple 平台上运行。示例用 coremltools 把 PyTorch 模型转成 Core ML 的 mlpackage

import coremltools as ct
import torch
import Colorizer

torch_model = Colorizer().eval()

example_input = torch.rand([1, 1, 256, 256])
traced_model = torch.jit.trace(torch_model, example_input)

coreml_model = ct.convert(traced_model, 
                          inputs=[ct.TensorType(name="input", shape=example_input.shape)])

coreml_model.save("Colorizer.mlpackage")

关键点:

  • Colorizer().eval() 加载模型架构和权重,并切到推理模式。
  • example_input 定义一个 1 通道、256×256 的输入样例。
  • torch.jit.trace 用样例输入 trace PyTorch 模型,得到可转换表示。
  • ct.convert 把 traced 模型转换成 Core ML 模型,并声明输入名称和形状。
  • save("Colorizer.mlpackage") 生成可以拖进 Xcode 的模型包。

3. 在 Python 里验证 Core ML 模型(04:26

转换成功不等于功能正确。演讲者在 Python 中直接运行 Core ML 模型,检查它和原 PyTorch 模型的流程是否一致。

import coremltools as ct
from PIL import Image
from skimage import color

in_img = Image.open(“image.png").convert("RGB")
in_rgb = np.array(in_img)
in_lab = color.rgb2lab(in_rgb, channel_axis=2)

lab_components = np.split(in_lab, indices_or_sections=3, axis=-1)
(in_l, _, _) = [
    np.expand_dims(array.transpose((2, 0, 1)).astype(np.float32), 0)
    for array in lab_components
]
out_ab = coreml_model.predict({"input": in_l})[0]

out_lab = np.squeeze(np.concatenate([in_l, out_ab], axis=1), axis=0).transpose((1, 2, 0))
out_rgb = color.lab2rgb(out_lab, channel_axis=2).astype(np.uint8)
out_img = Image.fromarray(out_rgb)

关键点:

  • Image.open(...).convert("RGB") 读取测试图片并统一为 RGB。
  • color.rgb2lab(..., channel_axis=2) 按最后一维作为颜色通道转换到 Lab。
  • np.split(..., axis=-1) 拆出 L、a、b 三个通道。
  • np.expand_dims(..., 0) 增加 batch 维,形成模型输入。
  • coreml_model.predict({"input": in_l}) 用转换后的 Core ML 模型预测颜色通道。
  • np.concatenate([in_l, out_ab], axis=1) 把输入亮度和输出颜色合并。
  • color.lab2rgb 转回 RGB,用 Image.fromarray 生成可查看图片。

4. 在 Swift App 中接入模型(07:11

Swift 里的集成流程和 Python 一样:提取亮度,送入 Core ML,取回两个颜色通道,再重建 RGB 图像。

import CoreImage
import CoreML

func colorize(image inputImage: CIImage) throws -> CIImage {

    let lightness: CIImage = extractLightness(from: inputImage)

    let modelInput = try ColorizerInput(inputWith: lightness.cgImage!)
    
    let modelOutput: ColorizerOutput = try colorizer.prediction(input: modelInput)

    let (aChannel, bChannel): (CIImage, CIImage) = extractColorChannels(from: modelOutput)

    let colorizedImage = reconstructRGBImage(l: lightness, a: aChannel, b: bChannel)
    return colorizedImage
}

关键点:

  • extractLightness(from:) 对应 Python 里的 RGB 到 Lab,再取 L 通道。
  • ColorizerInput(inputWith:) 把亮度图像准备成 Core ML 输入。
  • colorizer.prediction(input:) 执行模型推理。
  • extractColorChannels(from:) 把模型输出的两个 MLShapedArray 转成 Core Image 可处理的图像。
  • reconstructRGBImage 把 L、a、b 三个通道合成 Lab,再转回 RGB。

5. 用 Core Image 做 Lab 转换和通道合成(07:4108:51

演讲者没有在 Swift 里手写色彩空间转换。Core Image 提供了新的 convertRGBtoLabconvertLabToRGB 过滤器。

import CoreImage.CIFilterBuiltins

func extractLightness(from inputImage: CIImage) -> CIImage {

    let rgbToLabFilter = CIFilter.convertRGBtoLab()
    rgbToLabFilter.inputImage = inputImage
    rgbToLabFilter.normalize = true
    let labImage = rgbToLabFilter.outputImage

    let matrixFilter = CIFilter.colorMatrix()
    matrixFilter.inputImage = labImage
    matrixFilter.rVector = CIVector(x: 1, y: 0, z: 0)
    matrixFilter.gVector = CIVector(x: 1, y: 0, z: 0)
    matrixFilter.bVector = CIVector(x: 1, y: 0, z: 0)
    let lightness = matrixFilter.outputImage!
    return lightness
}

关键点:

  • CIFilter.convertRGBtoLab() 把输入图像从 RGB 转成 Lab。
  • normalize = true 让亮度值落在模型需要的归一化范围。
  • CIFilter.colorMatrix() 用矩阵筛出 L 通道。
  • rVectorgVectorbVector 都取输入的第一个分量,让输出图像只保留亮度。

模型输出是两个 MLShapedArray。示例把它们转成两个单通道 CIImage

func extractColorChannels(from output: ColorizerOutput) -> (CIImage, CIImage) {

    let outA: [Float] = output.output_aShapedArray.scalars
    let outB: [Float] = output.output_bShapedArray.scalars
    let dataA = Data(bytes: outA, count: outA.count * MemoryLayout<Float>.stride)
    let dataB = Data(bytes: outB, count: outB.count * MemoryLayout<Float>.stride)

    let outImageA = CIImage(bitmapData: dataA,
        bytesPerRow: 4 * 256,
        size: CGSize(width: 256, height: 256),
        format: CIFormat.Lh,
        colorSpace: CGColorSpaceCreateDeviceGray())
    let outImageB = CIImage(bitmapData: dataB,
        bytesPerRow: 4 * 256,
        size: CGSize(width: 256, height: 256),
        format: CIFormat.Lh,
        colorSpace: CGColorSpaceCreateDeviceGray())
   return (outImageA, outImageB)
}

关键点:

  • output_aShapedArray.scalarsoutput_bShapedArray.scalars 取出模型预测的颜色分量。
  • Data(bytes:count:) 把浮点数组包装成图像可读的字节数据。
  • CIImage(bitmapData:...) 为每个颜色分量创建单通道 Core Image 图像。
  • bytesPerRow: 4 * 256 对应 256 像素宽、每个 Float 4 字节。
  • CIFormat.Lh 表示单通道浮点图像格式。

最后用自定义 CIKernel 把三个通道合成 Lab 图像,再转回 RGB。

func reconstructRGBImage(l lightness: CIImage,
                         a aChannel: CIImage,
                         b bChannel: CIImage) -> CIImage {
    guard
        let kernel = try? CIKernel.kernels(withMetalString: source)[0] as? CIColorKernel,
        let kernelOutputImage = kernel.apply(extent: lightness.extent,
                                             arguments: [lightness, aChannel, bChannel])
    else { fatalError() }

    let labToRGBFilter = CIFilter.convertLabToRGBFilter()
    labToRGBFilter.inputImage = kernelOutputImage
    labToRGBFilter.normalize = true
    let rgbImage = labToRGBFilter.outputImage!
    return rgbImage
}

关键点:

  • CIKernel.kernels(withMetalString:) 从 Metal 字符串创建 Core Image kernel。
  • kernel.apply 以亮度图像范围为输出范围,把 L、a、b 作为参数传入。
  • CIFilter.convertLabToRGBFilter() 把合成后的 Lab 图像转回 RGB。
  • normalize = true 和前面的 RGB 到 Lab 转换保持配套。

通道合成本身只有一个很小的 kernel。

let source = """
#include <CoreImage/CoreImage.h>
[[stichable]] float4 labCombine(coreimage::sample_t imL, coreimage::sample_t imA, coreimage::sample_t imB)
{
   return float4(imL.r, imA.r, imB.r, imL.a);
}
"""

关键点:

  • imLimAimB 分别代表亮度、a 通道和 b 通道图像采样。
  • float4(imL.r, imA.r, imB.r, imL.a) 把三个单通道值放回 Lab 图像的前三个分量。
  • alpha 使用 imL.a,保持输入亮度图像的透明度。

6. 从 90ms 静态处理走向 16ms 实时处理(05:1712:45

Xcode 14 的 Core ML Performance Report 可以对 Core ML 模型做基于时间的分析。演讲者把模型拖进 Xcode,几秒内看到原模型在 iPad Pro M1 上约 90ms。这个速度适合静态照片。

实时摄像头需要更低延迟。演讲者改动模型架构并重新训练,使用 PyTorch on Metal 利用 Apple Silicon 加速训练。新模型的 Performance Report 显示,它完全运行在 Neural Engine 上,预测时间约 16ms。

Performance Report 只说明模型本身。App 运行后仍然不顺滑,Instruments 的 Core ML 模板显示预测请求在堆积:第二次预测在第一次还没完成时就提交给 Core ML,后续请求继续排队。演讲者还发现摄像头帧率误设为 60fps。

下面是概念性控制流,表达 Session 中的修复思路;不是官方代码片段。

// Conceptual pseudocode based on the runtime issue shown in Instruments.
if previousPredictionDidFinish {
    processNextCameraFrame()
}

cameraFrameRate = 30

关键点:

  • previousPredictionDidFinish 表示只在上一次预测完成后处理新帧。
  • processNextCameraFrame() 对应一次 Core ML 预测请求,避免请求排队。
  • cameraFrameRate = 30 对应演讲中把误设的 60fps 改回目标 30fps。
  • 这个修复让 Core ML 正确地一次向 Neural Engine 派发一个预测,实时效果才稳定。

核心启发

  • 做一个旧照片扫描上色工具:用 Vision 的 VNDetectRectangleRequest 找到取景画面里的照片区域,再把裁剪结果送入 Core ML Colorizer。Session 在 15:26 展示了这个组合。

  • 给现有图片 App 加一个“本地 ML 预检”流程:模型进 App 前,先用 coremltools 在 Python 里跑同一张输入图,确认 Core ML 输出和原模型流程一致,再进入 Swift 集成。

  • 把模型性能评估提前到原型阶段:设计功能时就把候选模型拖进 Xcode 14 Performance Report,记录目标设备的预测时间和计算单元分布,避免 UI 做完后才发现实时性不够。

  • 做一个摄像头实时滤镜的排队保护:摄像头每帧都来,但 Core ML 预测需要时间。用 Instruments 检查预测是否堆积,再把处理策略改成“上一帧完成后再接下一帧”。

  • 用 Apple Silicon 训练小型替代模型:如果开源模型太慢,保留同样输入输出形状,换更轻的架构,用 PyTorch on Metal 重新训练,再按转换、验证、集成的流程替换模型。

关联 Session

评论

GitHub Issues · utterances