如何将姿态检测YOLOv7模型转换为Core ML模型用于iOS设备部署

一、前期准备 #

• 准备好训练完成的YOLOv7姿态检测模型（.pt格式，官方预训练的yolov7-pose.pt也可直接使用）
• 安装依赖库：执行pip install torch torchvision coremltools完成环境配置

二、模型转换流程 #

1. 将YOLOv7模型转为ONNX格式

Core ML对ONNX模型的兼容性更好，先转ONNX再转Core ML是更稳妥的路径：

import torch
from models.yolov7 import Model
import yaml

# 加载模型配置与权重
with open('cfg/yolov7-pose.yaml') as f:
    cfg = yaml.safe_load(f)
# 姿态检测为单类（人体），所以nc=1
model = Model(cfg, ch=3, nc=1)
model.load_state_dict(torch.load('yolov7-pose.pt')['model'])
model.eval()

# 生成dummy输入并导出ONNX
dummy_input = torch.randn(1, 3, 640, 640)  # YOLOv7默认输入尺寸
torch.onnx.export(model, dummy_input, 'yolov7-pose.onnx', 
                  opset_version=12, input_names=['input'], output_names=['output'])

2. 转换为Core ML模型

将ONNX模型转为iOS可直接调用的.mlmodel格式，同时配置图像输入参数，方便iOS端直接传入UIImage：

import coremltools as ct

# 加载ONNX模型
onnx_model = ct.utils.load_spec('yolov7-pose.onnx')
# 配置输入为图像类型，设置归一化系数（YOLOv7输入需归一化到0-1）
coreml_model = ct.convert(onnx_model, 
                          inputs=[ct.ImageType(name='input', shape=(1, 3, 640, 640), scale=1/255.0)])
# 保存Core ML模型
coreml_model.save('YOLOv7Pose.mlmodel')

三、iOS端集成与推理 #

1. 导入模型到Xcode

将生成的YOLOv7Pose.mlmodel拖入Xcode项目，勾选对应target，Xcode会自动生成模型的Swift接口。

2. 姿态检测核心代码示例

import CoreML
import UIKit

class PoseDetector {
    private let model: YOLOv7Pose

    init() {
        do {
            model = try YOLOv7Pose(configuration: MLModelConfiguration())
        } catch {
            fatalError("模型加载失败：\(error)")
        }
    }

    /// 从UIImage中检测姿态关键点
    func detectPose(from image: UIImage) -> [PoseLandmark]? {
        guard let resizedImage = image.resize(to: CGSize(width: 640, height: 640)),
              let pixelBuffer = resizedImage.toCVPixelBuffer() else {
            return nil
        }
        do {
            let output = try model.prediction(input: pixelBuffer)
            return parsePoseOutput(output.output)
        } catch {
            print("推理错误：\(error)")
            return nil
        }
    }

    /// 解析模型输出为姿态关键点数组
    private func parsePoseOutput(_ output: MLMultiArray) -> [PoseLandmark] {
        var landmarks = [PoseLandmark]()
        // YOLOv7 Pose输出包含1个检测框 + 17个关键点的(x,y,置信度)，共1+17*3=52个元素
        for index in 0..<17 {
            let x = output[[0, 4 + index*3] as [NSNumber]].floatValue
            let y = output[[0, 5 + index*3] as [NSNumber]].floatValue
            let confidence = output[[0, 6 + index*3] as [NSNumber]].floatValue
            landmarks.append(PoseLandmark(x: x, y: y, confidence: confidence))
        }
        return landmarks
    }
}

// UIImage扩展：调整尺寸与转换为CVPixelBuffer
extension UIImage {
    func resize(to size: CGSize) -> UIImage? {
        UIGraphicsBeginImageContextWithOptions(size, false, UIScreen.main.scale)
        defer { UIGraphicsEndImageContext() }
        draw(in: CGRect(origin: .zero, size: size))
        return UIGraphicsGetImageFromCurrentImageContext()
    }

    func toCVPixelBuffer() -> CVPixelBuffer? {
        let attributes = [
            kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
            kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue
        ] as CFDictionary
        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                         Int(size.width),
                                         Int(size.height),
                                         kCVPixelFormatType_32ARGB,
                                         attributes,
                                         &pixelBuffer)
        guard status == kCVReturnSuccess, let buffer = pixelBuffer else {
            return nil
        }
        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }
        guard let context = CGContext(data: CVPixelBufferGetBaseAddress(buffer),
                                      width: Int(size.width),
                                      height: Int(size.height),
                                      bitsPerComponent: 8,
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(buffer),
                                      space: CGColorSpaceCreateDeviceRGB(),
                                      bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue),
              let cgImage = self.cgImage else {
            return nil
        }
        context.draw(cgImage, in: CGRect(origin: .zero, size: size))
        return buffer
    }
}

// 姿态关键点结构体
struct PoseLandmark {
    let x: Float
    let y: Float
    let confidence: Float
}

四、视频流可视化适配 #

你已经在Android和浏览器端实现了视频流的姿态可视化，iOS端可复用相同的绘制逻辑：

• 用AVFoundation捕获摄像头实时帧，每帧传入PoseDetector获取关键点
• 通过CAShapeLayer将关键点和骨骼连线绘制在预览层上，或使用Metal实现高性能渲染，保持跨端可视化效果一致

注意事项 #

• 模型输入尺寸必须与训练时一致（默认640x640），iOS端传入图像前必须调整尺寸
• 若转换时遇到算子不支持问题，可尝试调整opset_version（如11或13），或添加convert_to="mlprogram"参数（适配iOS 15+）
• 实时推理时建议开启Core ML的GPU加速，在MLModelConfiguration中设置computeUnits = .all

内容的提问来源于stack exchange，提问作者Arnav Sharma