Skip to content

Commit da2931e

Browse files
tucan9389tucan9389
authored
Merge pull request #60 from tucan9389/new-model/imgclsmob/simplepose
[PR] Support simple baseline pose estimation model provied by imgclsmob
2 parents a35166b + d0a1ce0 commit da2931e

13 files changed

Lines changed: 356 additions & 56 deletions

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,3 +91,6 @@ iOSInjectionProject/
9191

9292
Pods/*
9393
Podfile.lock
94+
95+
# TensorFlow Lite Model
96+
*.tflite

‎PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,22 @@
2727
712A7FD12426691B00B043F9 /* PEFMCPMPoseEstimator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 712A7FD02426691B00B043F9 /* PEFMCPMPoseEstimator.swift */; };
2828
712A7FD324266EC700B043F9 /* pefm_hourglass_v1.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 712A7FD224266EC700B043F9 /* pefm_hourglass_v1.tflite */; };
2929
7138DCCF242142FE0048E1D2 /* TFLiteFlatArray.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */; };
30+
713A0479247905AC00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */; };
31+
713A047A247905AC00F90C77 /* simplepose_resnet50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */; };
32+
713A047B247905AC00F90C77 /* alphapose_fastseresnet101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */; };
33+
713A047C247905AC00F90C77 /* simplepose_mobile_resnet18_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */; };
34+
713A047D247905AC00F90C77 /* simplepose_resnet152b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */; };
35+
713A047E247905AC00F90C77 /* simplepose_mobile_resnet50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */; };
36+
713A047F247905AC00F90C77 /* simplepose_resnet101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */; };
37+
713A0480247905AC00F90C77 /* simplepose_resneta152b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */; };
38+
713A0481247905AC00F90C77 /* simplepose_resneta50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */; };
39+
713A0482247905AC00F90C77 /* simplepose_resneta101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */; };
40+
713A0483247905AC00F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */; };
41+
713A0484247905AC00F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */; };
42+
713A0485247905AC00F90C77 /* simplepose_resnet18_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */; };
43+
713A0486247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */; };
44+
713A04882479071D00F90C77 /* IMGCLSPoseEstimator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */; };
45+
71A1BC1624793BC5004EA5D9 /* Argmax.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1BC1524793BC5004EA5D9 /* Argmax.swift */; };
3046
71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED1E24527D55001F796C /* PoseConfidenceMapDrawingView.swift */; };
3147
71A1ED4124574F2E001F796C /* StillImageHeatmapViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */; };
3248
71A1ED43245971AC001F796C /* LICENSE in Resources */ = {isa = PBXBuildFile; fileRef = 71A1ED42245971AC001F796C /* LICENSE */; };
@@ -64,6 +80,22 @@
6480
712A7FD02426691B00B043F9 /* PEFMCPMPoseEstimator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PEFMCPMPoseEstimator.swift; sourceTree = "<group>"; };
6581
712A7FD224266EC700B043F9 /* pefm_hourglass_v1.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = pefm_hourglass_v1.tflite; sourceTree = "<group>"; };
6682
7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TFLiteFlatArray.swift; sourceTree = "<group>"; };
83+
713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv3_small_w1_coco.tflite; sourceTree = "<group>"; };
84+
713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet50b_coco.tflite; sourceTree = "<group>"; };
85+
713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = alphapose_fastseresnet101b_coco.tflite; sourceTree = "<group>"; };
86+
713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_resnet18_coco.tflite; sourceTree = "<group>"; };
87+
713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet152b_coco.tflite; sourceTree = "<group>"; };
88+
713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_resnet50b_coco.tflite; sourceTree = "<group>"; };
89+
713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet101b_coco.tflite; sourceTree = "<group>"; };
90+
713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta152b_coco.tflite; sourceTree = "<group>"; };
91+
713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta50b_coco.tflite; sourceTree = "<group>"; };
92+
713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta101b_coco.tflite; sourceTree = "<group>"; };
93+
713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv2b_w1_coco.tflite; sourceTree = "<group>"; };
94+
713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv3_large_w1_coco.tflite; sourceTree = "<group>"; };
95+
713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet18_coco.tflite; sourceTree = "<group>"; };
96+
713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenet_w1_coco.tflite; sourceTree = "<group>"; };
97+
713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IMGCLSPoseEstimator.swift; sourceTree = "<group>"; };
98+
71A1BC1524793BC5004EA5D9 /* Argmax.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Argmax.swift; sourceTree = "<group>"; };
6799
71A1ED1E24527D55001F796C /* PoseConfidenceMapDrawingView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PoseConfidenceMapDrawingView.swift; sourceTree = "<group>"; };
68100
71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StillImageHeatmapViewController.swift; sourceTree = "<group>"; };
69101
71A1ED42245971AC001F796C /* LICENSE */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = LICENSE; sourceTree = "<group>"; };
@@ -154,6 +186,7 @@
154186
7105C92E241D0235001A4325 /* PoseEstimator.swift */,
155187
7105C932241D0651001A4325 /* TFLiteImageInterpreter.swift */,
156188
7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */,
189+
71FE62A62478EFB800425708 /* imgclsmob-Pose */,
157190
712A7FC22424FEC900B043F9 /* PoseNet */,
158191
712A7FC7242667A400B043F9 /* PoseEstimationForMobile */,
159192
71E8D9152438BA5B0081DD6E /* OpenPose */,
@@ -194,6 +227,7 @@
194227
isa = PBXGroup;
195228
children = (
196229
71DD577E2446D7CF0024C146 /* NonMaximumnonSuppression.swift */,
230+
71A1BC1524793BC5004EA5D9 /* Argmax.swift */,
197231
);
198232
name = Algorithm;
199233
sourceTree = "<group>";
@@ -207,6 +241,28 @@
207241
path = OpenPose;
208242
sourceTree = "<group>";
209243
};
244+
71FE62A62478EFB800425708 /* imgclsmob-Pose */ = {
245+
isa = PBXGroup;
246+
children = (
247+
713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */,
248+
713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */,
249+
713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */,
250+
713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */,
251+
713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */,
252+
713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */,
253+
713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */,
254+
713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */,
255+
713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */,
256+
713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */,
257+
713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */,
258+
713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */,
259+
713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */,
260+
713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */,
261+
713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */,
262+
);
263+
path = "imgclsmob-Pose";
264+
sourceTree = "<group>";
265+
};
210266
C0C61619494007101B3B8411 /* Pods */ = {
211267
isa = PBXGroup;
212268
children = (
@@ -284,15 +340,29 @@
284340
buildActionMask = 2147483647;
285341
files = (
286342
712A7FD324266EC700B043F9 /* pefm_hourglass_v1.tflite in Resources */,
343+
713A047D247905AC00F90C77 /* simplepose_resnet152b_coco.tflite in Resources */,
344+
713A0481247905AC00F90C77 /* simplepose_resneta50b_coco.tflite in Resources */,
287345
7105C922241CE9B7001A4325 /* LaunchScreen.storyboard in Resources */,
346+
713A047A247905AC00F90C77 /* simplepose_resnet50b_coco.tflite in Resources */,
288347
7105C91F241CE9B7001A4325 /* Assets.xcassets in Resources */,
348+
713A0479247905AC00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite in Resources */,
289349
712A7FCF2426690A00B043F9 /* pefm_cpm.tflite in Resources */,
350+
713A0486247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite in Resources */,
290351
71A1ED43245971AC001F796C /* LICENSE in Resources */,
352+
713A0483247905AC00F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite in Resources */,
291353
71E8D9172438BAC10081DD6E /* openpose_ildoonet.tflite in Resources */,
354+
713A0480247905AC00F90C77 /* simplepose_resneta152b_coco.tflite in Resources */,
292355
7105C91D241CE9B6001A4325 /* Main.storyboard in Resources */,
293356
71E8D93B243CC5330081DD6E /* README.md in Resources */,
357+
713A0482247905AC00F90C77 /* simplepose_resneta101b_coco.tflite in Resources */,
358+
713A0485247905AC00F90C77 /* simplepose_resnet18_coco.tflite in Resources */,
359+
713A047C247905AC00F90C77 /* simplepose_mobile_resnet18_coco.tflite in Resources */,
360+
713A047F247905AC00F90C77 /* simplepose_resnet101b_coco.tflite in Resources */,
294361
712A7FCD242668FF00B043F9 /* pefm_hourglass_v2.tflite in Resources */,
362+
713A047B247905AC00F90C77 /* alphapose_fastseresnet101b_coco.tflite in Resources */,
363+
713A0484247905AC00F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite in Resources */,
295364
7105C93A241E7624001A4325 /* posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite in Resources */,
365+
713A047E247905AC00F90C77 /* simplepose_mobile_resnet50b_coco.tflite in Resources */,
296366
);
297367
runOnlyForDeploymentPostprocessing = 0;
298368
};
@@ -330,6 +400,7 @@
330400
files = (
331401
7105C933241D0651001A4325 /* TFLiteImageInterpreter.swift in Sources */,
332402
71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */,
403+
713A04882479071D00F90C77 /* IMGCLSPoseEstimator.swift in Sources */,
333404
7105C93C241E8CE3001A4325 /* CVPixelBufferExtension.swift in Sources */,
334405
712A7FC9242667C900B043F9 /* PEFMHourglassPoseEstimator.swift in Sources */,
335406
71B07B97245E5C6C001FD385 /* NumericExtension.swift in Sources */,
@@ -339,6 +410,7 @@
339410
71FE62A52466941400425708 /* LiveLineHeatmapViewController.swift in Sources */,
340411
71E8D9192438BAD80081DD6E /* OpenPosePoseEstimator.swift in Sources */,
341412
71DD577F2446D7CF0024C146 /* NonMaximumnonSuppression.swift in Sources */,
413+
71A1BC1624793BC5004EA5D9 /* Argmax.swift in Sources */,
342414
712A7FC4242504EB00B043F9 /* PoseKeypointsDrawingView.swift in Sources */,
343415
7138DCCF242142FE0048E1D2 /* TFLiteFlatArray.swift in Sources */,
344416
7105C93E241E90C2001A4325 /* DataExtension.swift in Sources */,

‎PoseEstimation-TFLiteSwift/Argmax.swift‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/*
2+
* Copyright Doyoung Gwak 2020
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
//
18+
// Argmax.swift
19+
// PoseEstimation-TFLiteSwift
20+
//
21+
// Created by Doyoung Gwak on 2020/05/23.
22+
// Copyright © 2020 Doyoung Gwak. All rights reserved.
23+
//
24+
25+
import Foundation
26+
27+
extension TFLiteFlatArray where Element==Float32 {
28+
func argmax(_ heatmapIndex: Int) -> (row: Int, col: Int, val: Element) {
29+
var maxInfo = (row: 0, col: 0, val: self[heatmap: 0, 0, 0, heatmapIndex])
30+
let height = dimensions[1]
31+
let width = dimensions[2]
32+
for row in 0..<height {
33+
for col in 0..<width {
34+
if self[heatmap: 0, row, col, heatmapIndex] > maxInfo.val {
35+
maxInfo = (row: row, col: col, val: self[0, row, col, heatmapIndex])
36+
}
37+
}
38+
}
39+
return maxInfo
40+
}
41+
}

‎PoseEstimation-TFLiteSwift/LiveLineHeatmapViewController.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ class LiveLineHeatmapViewController: UIViewController {
122122
humanType: humanType)
123123
}
124124

125-
let poseEstimator: PoseEstimator = OpenPosePoseEstimator()
125+
let poseEstimator: PoseEstimator = IMGCLSPoseEstimator()
126126

127127
override func viewDidLoad() {
128128
super.viewDidLoad()

‎PoseEstimation-TFLiteSwift/NonMaximumnonSuppression.swift‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,19 @@
1+
/*
2+
* Copyright Doyoung Gwak 2020
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
117
//
218
// NonMaximumnonSuppression.swift
319
// PoseEstimation-TFLiteSwift

‎PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift‎

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -230,10 +230,10 @@ private extension PoseEstimationOutput {
230230

231231
switch postprocessOptions.humanType {
232232
case .singlePerson:
233-
// <#TODO#> - use partThreshold & don't use `convertToKeypoints` and `makeLines`
234-
let keypoints = convertToKeypoints(from: outputs)
235-
let lines = makeLines(with: keypoints)
236-
humans = [Human(keypoints: keypoints, lines: lines)]
233+
let human = parseSinglePerson(outputs,
234+
partIndex: postprocessOptions.bodyPart,
235+
partThreshold: postprocessOptions.partThreshold)
236+
humans = [human]
237237
case .multiPerson(let pairThreshold, let nmsFilterSize, let maxHumanNumber):
238238
humans = parseMultiHuman(outputs,
239239
partIndex: postprocessOptions.bodyPart,
@@ -244,20 +244,13 @@ private extension PoseEstimationOutput {
244244
}
245245
}
246246

247-
func convertToKeypoints(from outputs: [TFLiteFlatArray<Float32>]) -> [Keypoint] {
248-
let output = outputs[0] // openpose_ildoonet.tflite only use the first output
247+
func parseSinglePerson(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?) -> Human {
248+
// openpose_ildoonet.tflite only use the first output
249+
let output = outputs[0]
249250

250251
// get (col, row)s from heatmaps
251252
let keypointIndexInfos: [(row: Int, col: Int, val: Float32)] = (0..<OpenPosePoseEstimator.Output.ConfidenceMap.count).map { heatmapIndex in
252-
var maxInfo = (row: 0, col: 0, val: output[heatmap: 0, 0, 0, heatmapIndex])
253-
for row in 0..<OpenPosePoseEstimator.Output.ConfidenceMap.height {
254-
for col in 0..<OpenPosePoseEstimator.Output.ConfidenceMap.width {
255-
if output[heatmap: 0, row, col, heatmapIndex] > maxInfo.val {
256-
maxInfo = (row: row, col: col, val: output[0, row, col, heatmapIndex])
257-
}
258-
}
259-
}
260-
return maxInfo
253+
return output.argmax(heatmapIndex)
261254
}
262255

263256
// get points from (col, row)s and offsets
@@ -270,19 +263,25 @@ private extension PoseEstimationOutput {
270263
return (point: CGPoint(x: x, y: y), score: score)
271264
}
272265

273-
return keypointInfos.map { keypointInfo in Keypoint(position: keypointInfo.point, score: keypointInfo.score) }
274-
}
275-
276-
func makeLines(with keypoints: [Keypoint]) -> [Human.Line] {
266+
let keypoints: [Keypoint?] = keypointInfos
267+
.map { keypointInfo -> Keypoint? in Keypoint(position: keypointInfo.point, score: keypointInfo.score) }
268+
.map { keypointInfo -> Keypoint? in
269+
guard let score = keypointInfo?.score, let partThreshold = partThreshold else { return keypointInfo }
270+
return (score > partThreshold) ? keypointInfo : nil
271+
}
272+
273+
// lines
277274
var keypointWithBodyPart: [OpenPosePoseEstimator.Output.BodyPart: Keypoint] = [:]
278275
OpenPosePoseEstimator.Output.BodyPart.allCases.enumerated().forEach { (index, bodyPart) in
279276
keypointWithBodyPart[bodyPart] = keypoints[index]
280277
}
281-
return OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
278+
let lines: [Human.Line] = OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
282279
guard let fromKeypoint = keypointWithBodyPart[line.from],
283280
let toKeypoint = keypointWithBodyPart[line.to] else { return nil }
284281
return (from: fromKeypoint, to: toKeypoint)
285282
}
283+
284+
return Human(keypoints: keypoints, lines: lines)
286285
}
287286

288287
func parseMultiHuman(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?, pairThreshold: Float?, nmsFilterSize: Int, maxHumanNumber: Int?) -> [Human] {

0 commit comments

Comments
 (0)