以前ViTPoseをiOS用に実装した時ですが、Heatmapの前処理と後処理をC++とOpenCVに依存する形で実装しました。
その時にOpenCV依存を無くしてSwiftだけで実装するにはAffineTransformを自前で用意する必要があることを知りました。
AffineTransformについては、特に前知識もなく、手をつける為のとっかかりがなかったこともあり放置していたのですが、Swiftで書かれたコードを見つけたこともあり、試しにSwiftで実装しました。
Perspective transform from quadrilateral to quadrilateral in Swift
C++で書かれたHeatmap処理をSwiftに書き直す
元のC++のコード(Heatmap処理)はPaddlePaddleのコードを流用します。
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/deploy/lite/src/keypoint_postprocess.cc
作成したコード
https://github.com/otmb/TopDownPoseEstimation/blob/main/TopDownPoseEstimation/KeypointProcess.swift
リポジトリはこちら
cv::getAffineTransformをSwiftの関数に置き換える
難解な計算式もコードに落とすとこのようになる。
参照先: Perspective transform from quadrilateral to quadrilateral in Swift using SIMD for matrix operations
func affineTransform(from: double3x3, to: double3x3) -> double3x3? {
let invA = from.inverse
if invA.determinant.isNaN {
return nil
}
let M = to * invA
return M
}
func cgAffineTransform(from: double3x3, to: double3x3) -> CGAffineTransform? {
guard let M = affineTransform(from: from, to: to) else {
return nil
}
let (m1, m2, m3) = M.columns
return CGAffineTransform(a: m1.x, b: m1.y, c: m2.x, d: m2.y, tx: m3.x, ty: m3.y)
}
他に何か特筆すべき点を挙げてみる
今回のコードの実装で何か特筆すべき点を挙げるとしますと、PaddlePaddleのコードで一部ポインタを用いて実装しているコードがあったのですが、SwiftではAccelerateを用いたコードで処理してあげる方が高速でしたので、そのように実装しました。
PaddlePaddleのコード(C++)
void get_max_preds(std::vector<float>& heatmap,
std::vector<int>& dim,
std::vector<float>& preds,
std::vector<float>& maxvals,
int batchid,
int joint_idx) {
int num_joints = dim[1];
int width = dim[3];
std::vector<int> idx;
idx.resize(num_joints * 2);
for (int j = 0; j < dim[1]; j++) {
float* index = &(
heatmap[batchid * num_joints * dim[2] * dim[3] + j * dim[2] * dim[3]]);
float* end = index + dim[2] * dim[3];
float* max_dis = std::max_element(index, end);
auto max_id = std::distance(index, max_dis);
maxvals[j] = *max_dis;
if (*max_dis > 0) {
preds[j * 2] = static_cast<float>(max_id % width);
preds[j * 2 + 1] = static_cast<float>(max_id / width);
}
}
}
Swiftで書き直したコード
func getMaxCoords(heatmap: [Double]) -> [MaxCoord] {
let width = Double(heatmapWidth)
return (0..<keypointsNumber).map { j in
let idx = j * heatmapHeight * heatmapWidth
let end = idx + heatmapHeight * heatmapWidth
let (maxIdx, maxValue) = vDSP.indexOfMaximum(heatmap[idx..<end])
let coord = CGPoint(
x: Double(maxIdx).truncatingRemainder(dividingBy: width),
y: Double(maxIdx) / width)
return MaxCoord(coord: coord, maxval: maxValue)
}
}
その他: 描画処理
こちらもOpenCVでのpolygonを使った描画からCoreGraphicsの描画に切り替えました。
好みはありますがCoreGraphicsを使った方がSwiftだけで完結できるのでアプリサイズや導入のお作法的には良さそうです。
Before
After
描画のコードは下記を参考にしました。
Detecting human body poses in an image
追記
MS COCO val2017
今回作成したコードが正常に動いているかval2017で精度評価。
モバイルでこの精度が出るのは中々凄いのではないでしょうか。Vision Transformerのモデルはすごいですね。
| Models | AP |
|---|---|
| yolov7-tiny_fp16 + vitpose-b256x192_fp16.mlmodel | 0.589 |
| yolov7-tiny_fp16 + vitpose_s256x192_wholebody_fp16.mlmodel | 0.579 |
| yolov7-tiny_fp16 + vitpose_b256x192_wholebody_fp16.mlmodel | 0.600 |