iOS-TensorFlow Lite的多人姿势估计

最近在研究TensorFlow Lite ，发现Google 还真是区别对待呀。安卓亲儿子什么都要完善一些，不光是文档清晰详细，连demo也是更加完整。像iOS这端的多人姿势预测就没有写demo，因此只能靠自己一点一点摸索。自己也是好一阵捣鼓终于完成了多人姿势预测

1.首先我们需要多人姿势估计的模型

模型可以去TensorFlow的官网上下载，多人姿势模型如图我们得选择TFLite这个模块

截屏2022-06-29 下午5.26.51.png

2.解读模型要求

inputs：

视频或图像的帧，表示为动态形状的int32（对于TF.js）或uint8（对于TF Lite）shape：1xHxWx3，其中H和W需要是32的倍数，可以在运行时确定。通道顺序为RGB，值为[0，255]（注意：输入图像的大小控制着速度与精度之间的权衡，因此请选择最适合您的应用程序的值）

outputs：

返回float32类型的[1,6,56]
其中第一个参数是固定值不用管
第二个是可以检测到的最大数量 为6。
第三个表示预测的边界框/关键点位置和分数，其中17 * 3（51个数据）表示17个关键点 第一个是X坐标，第二个是Y坐标，第三个是置信分数---[y_0, x_0, s_0, y_1, x_1, s_1, …, y_16, x_16, s_16]对应[鼻子、左眼、右眼、左耳、右耳、左肩、右肩、左肘、右肘、左腕、右腕、左髋、右髋、左膝、右膝、左踝、右踝]。剩下的5个---[ymin、xmin、ymax、xmax、score]表示边界框的区域和总的置信分数

3.部分关键代码

解释器的创建和输入输出的创建

// 创建文件路径
     guard let modelPath = Bundle.main.path(
        forResource: fileInfo.name,
        ofType: "tflite")// 格式一定要tflite，如果不是说明模块下载错误
    else {
      // 失败处理
    }
// 解释器的自定义 （也可以省略）
    var options = Interpreter.Options()
    options.threadCount = 4
// 创建解释器 interpreter 需要设置为全局  var interpreter: Interpreter
    interpreter = try Interpreter(modelPath: modelPath, options: options, delegates: delegates) 
// [1xHxWx3] H和W需要是32的倍数  根据需求调整越大越精准
    try interpreter.resizeInput(at: 0, to: Tensor.Shape.init([1,256,256,3]))
    try interpreter.allocateTensors()
    inputTensor = try interpreter.input(at: 0)
    try interpreter.invoke()
    outputTensor = try interpreter.output(at: 0)

数据结构

import UIKit

// 性能检测
struct Times {
  var preprocessing: TimeInterval
  var inference: TimeInterval
  var postprocessing: TimeInterval
  var total: TimeInterval { preprocessing + inference + postprocessing }
}

/// 身体的关键点
enum BodyPart: String, CaseIterable {
  case nose = "nose"
  case leftEye = "left eye"
  case rightEye = "right eye"
  case leftEar = "left ear"
  case rightEar = "right ear"
  case leftShoulder = "left shoulder"
  case rightShoulder = "right shoulder"
  case leftElbow = "left elbow"
  case rightElbow = "right elbow"
  case leftWrist = "left wrist"
  case rightWrist = "right wrist"
  case leftHip = "left hip"
  case rightHip = "right hip"
  case leftKnee = "left knee"
  case rightKnee = "right knee"
  case leftAnkle = "left ankle"
  case rightAnkle = "right ankle"

  var position: Int {
    return BodyPart.allCases.firstIndex(of: self) ?? 0
  }
}

struct KeyPoint {
  var bodyPart: BodyPart = .nose
  var coordinate: CGPoint = .zero
/// 置信分数
  var score: Float32 = 0.0
}

struct Person {
  var keyPoints: [KeyPoint]
/// 总置信分数
  var score: Float32
}

对相机捕捉到的像素缓存进行处理输出Data

private func multiPreprocess(_ pixelBuffer: CVPixelBuffer) -> Data?{
        let sourcePixelFormat = CVPixelBufferGetPixelFormatType(pixelBuffer)
        assert(
          sourcePixelFormat == kCVPixelFormatType_32BGRA
            || sourcePixelFormat == kCVPixelFormatType_32ARGB)
       // 如果一开始没有resize shape这个地方输出的将是[1x1x1x3]
        let dimensions = inputTensor.shape.dimensions
        let inputWidth = CGFloat(dimensions[1])// 32的倍数才行
        let inputHeight = CGFloat(dimensions[2])// 32的倍数才行
        let imageWidth = CGFloat(Int(pixelBuffer.size.width/32) * 32)
        let imageHeight = CGFloat(Int(pixelBuffer.size.height/32) * 32)

        let cropRegion = self.cropRegion ??
          initialCropRegion(imageWidth: imageWidth, imageHeight: imageHeight)
        self.cropRegion = cropRegion

        let rectF = RectF(
          left: cropRegion.left * imageWidth,
          top: cropRegion.top * imageHeight,
          right: cropRegion.right * imageWidth,
          bottom: cropRegion.bottom * imageHeight)

        // Detect region
        let modelSize = CGSize(width: inputWidth, height: inputHeight)
// 处理图像数据可以用TensorFlow example中的方法
        guard let thumbnail = pixelBuffer.cropAndResize(fromRect: rectF.rect, toSize: modelSize) else {
          return nil
        }

        // Remove the alpha component from the image buffer to get the initialized `Data`.
        guard
          let inputData = thumbnail.rgbData(
            isModelQuantized: inputTensor.dataType == .uInt8, imageMean: imageMean, imageStd: imageStd)
        else {
          os_log("Failed to convert the image buffer to RGB data.", type: .error)
          return nil
        }
        return inputData
    }

输出模型

private func postMultiprocess(imageSize: CGSize, modelOutput: Tensor) -> [Person]?{
        // 长宽处理成32的倍数
        let imageWidth = CGFloat(Int(imageSize.width/32) * 32)
        let imageHeight = CGFloat(Int(imageSize.height/32) * 32)

        let cropRegion = self.cropRegion ??
          initialCropRegion(imageWidth: imageWidth, imageHeight: imageHeight)

        let minX: CGFloat = cropRegion.left * imageWidth
        let minY: CGFloat = cropRegion.top * imageHeight

        var output = modelOutput.data.toArray(type: Float32.self)
        var dataArr = Array<Array<Float32>>()
        let dimensions = modelOutput.shape.dimensions
        let numKeyPoints = dimensions[2]
          while output.count > 0 {
              var tempArr = [Float32]()
              for (indexNum, obj) in output.enumerated() {
                  if indexNum < numKeyPoints{
                      tempArr.append(obj)
                      output.remove(at: 0)
                  }
              }
              dataArr.append(tempArr)
          }
        // 批处理维度，它总是等于1
        _ = CGFloat(modelOutput.shape.dimensions[1])
        // 实例检测的最大数量(6)
        _ = CGFloat(modelOutput.shape.dimensions[2])
        
        let inputWidth = CGFloat(inputTensor.shape.dimensions[1])
        let inputHeight = CGFloat(inputTensor.shape.dimensions[2])
        let widthRatio = (cropRegion.width * imageWidth / inputWidth)
        let heightRatio = (cropRegion.height * imageHeight / inputHeight)

        // Translate the coordinates from the model output's [0..1] back to that of
        // the input image
        
        var personArr = [Person]()
        for arr in dataArr {
            var keyPoints: [KeyPoint] = []
            for idx in 0...16{
                let x = ((CGFloat(arr[idx * 3 + 1]) * inputWidth) * widthRatio) + minX
                let y = ((CGFloat(arr[idx * 3 + 0]) * inputHeight) * heightRatio) + minY
                let score = arr[idx * 3 + 2]
                let keyPoint = KeyPoint(
                  bodyPart: BodyPart.allCases[idx], coordinate: CGPoint(x: x, y: y), score: score)
                keyPoints.append(keyPoint)
            }
            let totalScore = arr[numKeyPoints - 1]
            personArr.append(Person(keyPoints: keyPoints, score: totalScore))
        }

        return personArr
    }

多人姿势估计

func estimateMultiPose(on pixelBuffer: CVPixelBuffer) throws -> ([Person], Times) {
        guard !isProcessing else {
          throw PoseEstimationError.modelBusy
        }
        isProcessing = true
        defer {
          isProcessing = false
        }
        // 每次肢体检测的开始时间
        let preprocessingStartTime: Date
        let inferenceStartTime: Date
        let postprocessingStartTime: Date

        // 过程的时间
        let preprocessingTime: TimeInterval
        let inferenceTime: TimeInterval
        let postprocessingTime: TimeInterval

        preprocessingStartTime = Date()
        guard let data = multiPreprocess(pixelBuffer) else {
          throw PoseEstimationError.preprocessingFailed
        }
        preprocessingTime = Date().timeIntervalSince(preprocessingStartTime)

        inferenceStartTime = Date()
        do {
          // 将数据拷贝进解释器
          try interpreter.copy(data, toInputAt: 0)
          try interpreter.invoke()
          outputTensor = try interpreter.output(at: 0)
        } catch {
          throw PoseEstimationError.inferenceFailed
        }
        inferenceTime = Date().timeIntervalSince(inferenceStartTime)
        postprocessingStartTime = Date()
        guard let resultArr = postMultiprocess(imageSize: pixelBuffer.size, modelOutput: outputTensor) else {
          throw PoseEstimationError.postProcessingFailed
        }
        postprocessingTime = Date().timeIntervalSince(postprocessingStartTime)

        let times = Times(
          preprocessing: preprocessingTime,
          inference: inferenceTime,
          postprocessing: postprocessingTime)
        return (resultArr, times)
    }

到此整个处理的流程就结束了，CVPixelBuffer可以通过相机捕获

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

也可以自己找图片转换（CGImage转换成CVPixelBuffer）。反正很灵活，想怎么玩都行