自定义数据集训练的YOLOv8s转CoreML后，如何获取检测框坐标？

我用自定义数据集训练了YOLOv8目标检测模型，导出为Core ML模型后，需要获取检测目标的坐标/边界框，才能在iOS上给目标绘制矩形框，作为初学者不知道怎么实现，求帮助。

训练、验证、导出代码如下：

训练代码 #

!yolo task=detect mode=train model=yolov8s.pt data=data.yaml epochs=25 imgsz=640 plots=True

验证代码 #

!yolo task=detect mode=val model=runs/detect/train/weights/best.pt data=data.yaml

导出Core ML代码 #

!yolo mode=export model=runs/detect/train/weights/best.pt format=coreml

1. 调整Core ML导出参数，确保输出可解析的检测框 #

默认导出的Core ML模型可能输出格式不便于直接提取坐标，你可以添加参数让模型直接输出过滤后的检测框数据：

!yolo mode=export model=runs/detect/train/weights/best.pt format=coreml imgsz=640 nms=True

• nms=True：开启非极大值抑制，自动过滤重复、低质量的检测框
• imgsz=640：和训练时保持一致，避免输入输出尺寸不匹配导致坐标偏移

导出后的模型输出会包含核心字段：

• confidence：检测结果的置信度
• boundingBox：归一化的边界框坐标（格式为x_center, y_center, width, height，数值范围0-1）
• classLabel：检测到的目标类别

2. 在iOS端解析模型输出并转换为可绘制的坐标 #

步骤1：导入模型到iOS项目 #

将导出的.mlmodel文件拖入Xcode项目，Xcode会自动生成对应的模型类（比如Best）。

步骤2：预处理输入图像 #

YOLOv8要求输入为640x640尺寸，且像素值归一化到0-1范围，需要将iOS端的图像转换为符合要求的格式：

import CoreML
import Vision

// 自定义图像转换扩展
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 attrs = [
            kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
            kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue
        ] as CFDictionary
        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(
            kCFAllocatorDefault,
            Int(size.width),
            Int(size.height),
            kCVPixelFormatType_32BGRA,
            attrs,
            &pixelBuffer
        )
        guard status == kCVReturnSuccess, let pb = pixelBuffer else { return nil }
        
        CVPixelBufferLockBaseAddress(pb, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pb, .readOnly) }
        let context = CGContext(
            data: CVPixelBufferGetBaseAddress(pb),
            width: Int(size.width),
            height: Int(size.height),
            bitsPerComponent: 8,
            bytesPerRow: CVPixelBufferGetBytesPerRow(pb),
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.premultipliedFirst.rawValue
        )
        context?.draw(cgImage!, in: CGRect(origin: .zero, size: size))
        return pb
    }
}

// 预处理图像
func preprocessImage(_ image: UIImage) -> CVPixelBuffer? {
    guard let resizedImage = image.resize(to: CGSize(width: 640, height: 640)),
          let pixelBuffer = resizedImage.toCVPixelBuffer() else {
        return nil
    }
    // 归一化像素值到0-1范围
    CVPixelBufferLockBaseAddress(pixelBuffer, .readWrite)
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readWrite) }
    let floatBuffer = unsafeBitCast(CVPixelBufferGetBaseAddress(pixelBuffer), to: UnsafeMutablePointer<Float>.self)
    let pixelCount = 640 * 640 * 4
    for i in 0..<pixelCount {
        floatBuffer[i] = floatBuffer[i] / 255.0
    }
    return pixelBuffer
}

步骤3：运行模型并解析检测框坐标 #

// 加载模型
guard let model = try? Best(configuration: MLModelConfiguration()) else {
    fatalError("无法加载Core ML模型")
}

// 创建Vision检测请求
let request = VNCoreMLRequest(model: VNCoreMLModel(for: model.model)) { request, error in
    guard let results = request.results as? [VNRecognizedObjectObservation] else {
        return
    }
    
    // 遍历检测结果，转换为可绘制的坐标
    let originalImageSize = yourOriginalImage.size // 替换为你的原始图像尺寸
    for result in results {
        let confidence = result.confidence
        guard confidence >= 0.5 else { // 过滤低置信度结果
            continue
        }
        
        // Vision返回的是归一化坐标，原点在图像左下角，需转换为iOS UI坐标系（左上角为原点）
        let normalizedBox = result.boundingBox
        let x = normalizedBox.origin.x * originalImageSize.width
        let y = (1 - normalizedBox.origin.y - normalizedBox.size.height) * originalImageSize.height
        let width = normalizedBox.size.width * originalImageSize.width
        let height = normalizedBox.size.height * originalImageSize.height
        
        let drawRect = CGRect(x: x, y: y, width: width, height: height)
        // 此处可使用drawRect在图像上绘制矩形框
        print("目标坐标：\(drawRect)，置信度：\(confidence)")
    }
}

// 执行检测请求
if let pixelBuffer = preprocessImage(yourOriginalImage) {
    let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer)
    try? handler.perform([request])
}

关键注意事项 #

• 坐标系转换：Vision框架的坐标原点在图像左下角，iOS UI坐标系原点在左上角，必须对y轴进行转换，否则绘制的框会位置颠倒。
• 置信度阈值：根据实际需求调整阈值（比如0.5），过滤掉不可靠的检测结果。
• 尺寸匹配：输入图像尺寸必须和训练、导出时的imgsz一致，否则坐标会出现比例错误。

内容的提问来源于stack exchange，提问作者Dhiman Das