独立开发者

Hackaton：使用PoseNet模型检测图像中人体姿态

编辑

跳转Github地址

预估时长 1 min read

项目名称：Hackaton

Hackaton项目是一个开源的Swift项目，它展示了如何使用第三方Core ML模型PoseNet来从摄像头捕获的图像帧中检测人体姿态。PoseNet模型能够检测17个不同的身体部位或关节，包括眼睛、耳朵、鼻子、肩膀、臀部、肘部、膝盖、手腕和脚踝。这些关节共同构成了一个人的姿态。

项目的主要功能包括：

配置捕获会话：使用AVCaptureSession从设备内置摄像头获取图像。
获取捕获的图像：将视频捕获会话发送的每张图像转换为CGImage，然后传递给VideoCapture对象的代理。
准备PoseNet模型的输入：将捕获的图像包装在PoseNetInput实例中，调整图像大小以符合模型输入要求。
将输入传递给PoseNet模型：使用PoseNet模型的prediction函数获取输出，用于检测姿态。
分析PoseNet输出以定位关节：使用算法定位单个或多个人的姿态关节。
可视化检测到的姿态：在输入图像上绘制线框，连接关节，并绘制关节本身。

以下是项目中的关键代码段，展示了如何实现这些功能：

// 配置捕获会话
captureSession.beginConfiguration()
captureSession.sessionPreset = .vga640x480
try setCaptureSessionInput()
try setCaptureSessionOutput()
captureSession.commitConfiguration()

// 获取捕获的图像
guard CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly) == kCVReturnSuccess else {
    return
}
var image: CGImage?
VTCreateCGImageFromCVPixelBuffer(pixelBuffer, options: nil, imageOut: &image)
CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
DispatchQueue.main.sync {
    delegate.videoCapture(self, didCaptureFrame: image)
}

// 准备PoseNet模型的输入
let input = PoseNetInput(image: image, size: self.modelInputSize)

// 将输入传递给PoseNet模型
guard let prediction = try? self.poseNetMLModel.prediction(from: input) else {
    return
}
let poseNetOutput = PoseNetOutput(prediction: prediction, modelInputSize: self.modelInputSize, modelOutputStride: self.outputStride)
DispatchQueue.main.async {
    self.delegate?.poseNet(self, didPredict: poseNetOutput)
}

// 分析PoseNet输出以定位关节
var pose = Pose()
pose.joints.values.forEach { joint in
    configure(joint: joint)
}
pose.confidence = pose.joints.values.map { $0.confidence }.reduce(0, +) / Double(Joint.numberOfJoints)
pose.joints.values.forEach { joint in
    joint.position = joint.position.applying(modelToInputTransformation)
}

// 可视化检测到的姿态
let dstImageSize = CGSize(width: frame.width, height: frame.height)
let dstImageFormat = UIGraphicsImageRendererFormat()
dstImageFormat.scale = 1
let renderer = UIGraphicsImageRenderer(size: dstImageSize, format: dstImageFormat)
let dstImage = renderer.image { rendererContext in
    draw(image: frame, in: rendererContext.cgContext)
    for pose in poses {
        for segment in PoseImageView.jointSegments {
            let jointA = pose[segment.jointA]
            let jointB = pose[segment.jointB]
            guard jointA.isValid, jointB.isValid else {
                continue
            }
            drawLine(from: jointA, to: jointB, in: rendererContext.cgContext)
        }
        for joint in pose.joints.values.filter({ $0.isValid }) {
            draw(circle: joint, in: rendererContext.cgContext)
        }
    }
}