Skip to content

Commit

Permalink
Merge pull request #60 from tucan9389/new-model/imgclsmob/simplepose
Browse files Browse the repository at this point in the history
[PR] Support simple baseline pose estimation model provied by imgclsmob
  • Loading branch information
tucan9389 committed May 24, 2020
2 parents a35166b + d0a1ce0 commit da2931e
Show file tree
Hide file tree
Showing 13 changed files with 356 additions and 56 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -91,3 +91,6 @@ iOSInjectionProject/

Pods/*
Podfile.lock

# TensorFlow Lite Model
*.tflite
72 changes: 72 additions & 0 deletions PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,22 @@
712A7FD12426691B00B043F9 /* PEFMCPMPoseEstimator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 712A7FD02426691B00B043F9 /* PEFMCPMPoseEstimator.swift */; };
712A7FD324266EC700B043F9 /* pefm_hourglass_v1.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 712A7FD224266EC700B043F9 /* pefm_hourglass_v1.tflite */; };
7138DCCF242142FE0048E1D2 /* TFLiteFlatArray.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */; };
713A0479247905AC00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */; };
713A047A247905AC00F90C77 /* simplepose_resnet50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */; };
713A047B247905AC00F90C77 /* alphapose_fastseresnet101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */; };
713A047C247905AC00F90C77 /* simplepose_mobile_resnet18_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */; };
713A047D247905AC00F90C77 /* simplepose_resnet152b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */; };
713A047E247905AC00F90C77 /* simplepose_mobile_resnet50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */; };
713A047F247905AC00F90C77 /* simplepose_resnet101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */; };
713A0480247905AC00F90C77 /* simplepose_resneta152b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */; };
713A0481247905AC00F90C77 /* simplepose_resneta50b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */; };
713A0482247905AC00F90C77 /* simplepose_resneta101b_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */; };
713A0483247905AC00F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */; };
713A0484247905AC00F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */; };
713A0485247905AC00F90C77 /* simplepose_resnet18_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */; };
713A0486247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */; };
713A04882479071D00F90C77 /* IMGCLSPoseEstimator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */; };
71A1BC1624793BC5004EA5D9 /* Argmax.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1BC1524793BC5004EA5D9 /* Argmax.swift */; };
71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED1E24527D55001F796C /* PoseConfidenceMapDrawingView.swift */; };
71A1ED4124574F2E001F796C /* StillImageHeatmapViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */; };
71A1ED43245971AC001F796C /* LICENSE in Resources */ = {isa = PBXBuildFile; fileRef = 71A1ED42245971AC001F796C /* LICENSE */; };
Expand Down Expand Up @@ -64,6 +80,22 @@
712A7FD02426691B00B043F9 /* PEFMCPMPoseEstimator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PEFMCPMPoseEstimator.swift; sourceTree = "<group>"; };
712A7FD224266EC700B043F9 /* pefm_hourglass_v1.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = pefm_hourglass_v1.tflite; sourceTree = "<group>"; };
7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TFLiteFlatArray.swift; sourceTree = "<group>"; };
713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv3_small_w1_coco.tflite; sourceTree = "<group>"; };
713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet50b_coco.tflite; sourceTree = "<group>"; };
713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = alphapose_fastseresnet101b_coco.tflite; sourceTree = "<group>"; };
713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_resnet18_coco.tflite; sourceTree = "<group>"; };
713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet152b_coco.tflite; sourceTree = "<group>"; };
713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_resnet50b_coco.tflite; sourceTree = "<group>"; };
713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet101b_coco.tflite; sourceTree = "<group>"; };
713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta152b_coco.tflite; sourceTree = "<group>"; };
713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta50b_coco.tflite; sourceTree = "<group>"; };
713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resneta101b_coco.tflite; sourceTree = "<group>"; };
713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv2b_w1_coco.tflite; sourceTree = "<group>"; };
713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenetv3_large_w1_coco.tflite; sourceTree = "<group>"; };
713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_resnet18_coco.tflite; sourceTree = "<group>"; };
713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = simplepose_mobile_mobilenet_w1_coco.tflite; sourceTree = "<group>"; };
713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IMGCLSPoseEstimator.swift; sourceTree = "<group>"; };
71A1BC1524793BC5004EA5D9 /* Argmax.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Argmax.swift; sourceTree = "<group>"; };
71A1ED1E24527D55001F796C /* PoseConfidenceMapDrawingView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PoseConfidenceMapDrawingView.swift; sourceTree = "<group>"; };
71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StillImageHeatmapViewController.swift; sourceTree = "<group>"; };
71A1ED42245971AC001F796C /* LICENSE */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = LICENSE; sourceTree = "<group>"; };
Expand Down Expand Up @@ -154,6 +186,7 @@
7105C92E241D0235001A4325 /* PoseEstimator.swift */,
7105C932241D0651001A4325 /* TFLiteImageInterpreter.swift */,
7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */,
71FE62A62478EFB800425708 /* imgclsmob-Pose */,
712A7FC22424FEC900B043F9 /* PoseNet */,
712A7FC7242667A400B043F9 /* PoseEstimationForMobile */,
71E8D9152438BA5B0081DD6E /* OpenPose */,
Expand Down Expand Up @@ -194,6 +227,7 @@
isa = PBXGroup;
children = (
71DD577E2446D7CF0024C146 /* NonMaximumnonSuppression.swift */,
71A1BC1524793BC5004EA5D9 /* Argmax.swift */,
);
name = Algorithm;
sourceTree = "<group>";
Expand All @@ -207,6 +241,28 @@
path = OpenPose;
sourceTree = "<group>";
};
71FE62A62478EFB800425708 /* imgclsmob-Pose */ = {
isa = PBXGroup;
children = (
713A046D2479057B00F90C77 /* alphapose_fastseresnet101b_coco.tflite */,
713A0478247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite */,
713A0475247905A800F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite */,
713A046B2479056E00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite */,
713A0476247905A900F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite */,
713A046E2479057D00F90C77 /* simplepose_mobile_resnet18_coco.tflite */,
713A04702479058A00F90C77 /* simplepose_mobile_resnet50b_coco.tflite */,
713A0477247905AB00F90C77 /* simplepose_resnet18_coco.tflite */,
713A046C2479057300F90C77 /* simplepose_resnet50b_coco.tflite */,
713A04712479059200F90C77 /* simplepose_resnet101b_coco.tflite */,
713A046F2479058600F90C77 /* simplepose_resnet152b_coco.tflite */,
713A0473247905A000F90C77 /* simplepose_resneta50b_coco.tflite */,
713A0474247905A700F90C77 /* simplepose_resneta101b_coco.tflite */,
713A04722479059B00F90C77 /* simplepose_resneta152b_coco.tflite */,
713A04872479071D00F90C77 /* IMGCLSPoseEstimator.swift */,
);
path = "imgclsmob-Pose";
sourceTree = "<group>";
};
C0C61619494007101B3B8411 /* Pods */ = {
isa = PBXGroup;
children = (
Expand Down Expand Up @@ -284,15 +340,29 @@
buildActionMask = 2147483647;
files = (
712A7FD324266EC700B043F9 /* pefm_hourglass_v1.tflite in Resources */,
713A047D247905AC00F90C77 /* simplepose_resnet152b_coco.tflite in Resources */,
713A0481247905AC00F90C77 /* simplepose_resneta50b_coco.tflite in Resources */,
7105C922241CE9B7001A4325 /* LaunchScreen.storyboard in Resources */,
713A047A247905AC00F90C77 /* simplepose_resnet50b_coco.tflite in Resources */,
7105C91F241CE9B7001A4325 /* Assets.xcassets in Resources */,
713A0479247905AC00F90C77 /* simplepose_mobile_mobilenetv3_small_w1_coco.tflite in Resources */,
712A7FCF2426690A00B043F9 /* pefm_cpm.tflite in Resources */,
713A0486247905AC00F90C77 /* simplepose_mobile_mobilenet_w1_coco.tflite in Resources */,
71A1ED43245971AC001F796C /* LICENSE in Resources */,
713A0483247905AC00F90C77 /* simplepose_mobile_mobilenetv2b_w1_coco.tflite in Resources */,
71E8D9172438BAC10081DD6E /* openpose_ildoonet.tflite in Resources */,
713A0480247905AC00F90C77 /* simplepose_resneta152b_coco.tflite in Resources */,
7105C91D241CE9B6001A4325 /* Main.storyboard in Resources */,
71E8D93B243CC5330081DD6E /* README.md in Resources */,
713A0482247905AC00F90C77 /* simplepose_resneta101b_coco.tflite in Resources */,
713A0485247905AC00F90C77 /* simplepose_resnet18_coco.tflite in Resources */,
713A047C247905AC00F90C77 /* simplepose_mobile_resnet18_coco.tflite in Resources */,
713A047F247905AC00F90C77 /* simplepose_resnet101b_coco.tflite in Resources */,
712A7FCD242668FF00B043F9 /* pefm_hourglass_v2.tflite in Resources */,
713A047B247905AC00F90C77 /* alphapose_fastseresnet101b_coco.tflite in Resources */,
713A0484247905AC00F90C77 /* simplepose_mobile_mobilenetv3_large_w1_coco.tflite in Resources */,
7105C93A241E7624001A4325 /* posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite in Resources */,
713A047E247905AC00F90C77 /* simplepose_mobile_resnet50b_coco.tflite in Resources */,
);
runOnlyForDeploymentPostprocessing = 0;
};
Expand Down Expand Up @@ -330,6 +400,7 @@
files = (
7105C933241D0651001A4325 /* TFLiteImageInterpreter.swift in Sources */,
71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */,
713A04882479071D00F90C77 /* IMGCLSPoseEstimator.swift in Sources */,
7105C93C241E8CE3001A4325 /* CVPixelBufferExtension.swift in Sources */,
712A7FC9242667C900B043F9 /* PEFMHourglassPoseEstimator.swift in Sources */,
71B07B97245E5C6C001FD385 /* NumericExtension.swift in Sources */,
Expand All @@ -339,6 +410,7 @@
71FE62A52466941400425708 /* LiveLineHeatmapViewController.swift in Sources */,
71E8D9192438BAD80081DD6E /* OpenPosePoseEstimator.swift in Sources */,
71DD577F2446D7CF0024C146 /* NonMaximumnonSuppression.swift in Sources */,
71A1BC1624793BC5004EA5D9 /* Argmax.swift in Sources */,
712A7FC4242504EB00B043F9 /* PoseKeypointsDrawingView.swift in Sources */,
7138DCCF242142FE0048E1D2 /* TFLiteFlatArray.swift in Sources */,
7105C93E241E90C2001A4325 /* DataExtension.swift in Sources */,
Expand Down
41 changes: 41 additions & 0 deletions PoseEstimation-TFLiteSwift/Argmax.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
* Copyright Doyoung Gwak 2020
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

//
// Argmax.swift
// PoseEstimation-TFLiteSwift
//
// Created by Doyoung Gwak on 2020/05/23.
// Copyright © 2020 Doyoung Gwak. All rights reserved.
//

import Foundation

extension TFLiteFlatArray where Element==Float32 {
func argmax(_ heatmapIndex: Int) -> (row: Int, col: Int, val: Element) {
var maxInfo = (row: 0, col: 0, val: self[heatmap: 0, 0, 0, heatmapIndex])
let height = dimensions[1]
let width = dimensions[2]
for row in 0..<height {
for col in 0..<width {
if self[heatmap: 0, row, col, heatmapIndex] > maxInfo.val {
maxInfo = (row: row, col: col, val: self[0, row, col, heatmapIndex])
}
}
}
return maxInfo
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ class LiveLineHeatmapViewController: UIViewController {
humanType: humanType)
}

let poseEstimator: PoseEstimator = OpenPosePoseEstimator()
let poseEstimator: PoseEstimator = IMGCLSPoseEstimator()

override func viewDidLoad() {
super.viewDidLoad()
Expand Down
16 changes: 16 additions & 0 deletions PoseEstimation-TFLiteSwift/NonMaximumnonSuppression.swift
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
/*
* Copyright Doyoung Gwak 2020
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

//
// NonMaximumnonSuppression.swift
// PoseEstimation-TFLiteSwift
Expand Down
39 changes: 19 additions & 20 deletions PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift
Original file line number Diff line number Diff line change
Expand Up @@ -230,10 +230,10 @@ private extension PoseEstimationOutput {

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

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

// get (col, row)s from heatmaps
let keypointIndexInfos: [(row: Int, col: Int, val: Float32)] = (0..<OpenPosePoseEstimator.Output.ConfidenceMap.count).map { heatmapIndex in
var maxInfo = (row: 0, col: 0, val: output[heatmap: 0, 0, 0, heatmapIndex])
for row in 0..<OpenPosePoseEstimator.Output.ConfidenceMap.height {
for col in 0..<OpenPosePoseEstimator.Output.ConfidenceMap.width {
if output[heatmap: 0, row, col, heatmapIndex] > maxInfo.val {
maxInfo = (row: row, col: col, val: output[0, row, col, heatmapIndex])
}
}
}
return maxInfo
return output.argmax(heatmapIndex)
}

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

return keypointInfos.map { keypointInfo in Keypoint(position: keypointInfo.point, score: keypointInfo.score) }
}

func makeLines(with keypoints: [Keypoint]) -> [Human.Line] {
let keypoints: [Keypoint?] = keypointInfos
.map { keypointInfo -> Keypoint? in Keypoint(position: keypointInfo.point, score: keypointInfo.score) }
.map { keypointInfo -> Keypoint? in
guard let score = keypointInfo?.score, let partThreshold = partThreshold else { return keypointInfo }
return (score > partThreshold) ? keypointInfo : nil
}

// lines
var keypointWithBodyPart: [OpenPosePoseEstimator.Output.BodyPart: Keypoint] = [:]
OpenPosePoseEstimator.Output.BodyPart.allCases.enumerated().forEach { (index, bodyPart) in
keypointWithBodyPart[bodyPart] = keypoints[index]
}
return OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
let lines: [Human.Line] = OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
guard let fromKeypoint = keypointWithBodyPart[line.from],
let toKeypoint = keypointWithBodyPart[line.to] else { return nil }
return (from: fromKeypoint, to: toKeypoint)
}

return Human(keypoints: keypoints, lines: lines)
}

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

0 comments on commit da2931e

Please sign in to comment.