Skip to content

Commit

Permalink
Support simple baseline pose estimation model provied by imgclsmob re…
Browse files Browse the repository at this point in the history 
…po (#59)
  • Loading branch information
@tucan9389
tucan9389 committed May 23, 2020
1 parent a35166b commit 5dae993
Show file tree
Hide file tree
Showing 7 changed files with 297 additions and 14 deletions.
3 changes: 3 additions & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,3 +91,6 @@ iOSInjectionProject/

Pods/*
Podfile.lock

# TensorFlow Lite Model
*.tflite
68 changes: 68 additions & 0 deletions PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,21 @@
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 */; };
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 +79,21 @@
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>"; };
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 +184,7 @@
7105C92E241D0235001A4325 /* PoseEstimator.swift */,
7105C932241D0651001A4325 /* TFLiteImageInterpreter.swift */,
7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */,
71FE62A62478EFB800425708 /* imgclsmob-Pose */,
712A7FC22424FEC900B043F9 /* PoseNet */,
712A7FC7242667A400B043F9 /* PoseEstimationForMobile */,
71E8D9152438BA5B0081DD6E /* OpenPose */,
Expand Down Expand Up @@ -207,6 +238,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 +337,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 +397,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 Down
2 changes: 1 addition & 1 deletion PoseEstimation-TFLiteSwift/LiveLineHeatmapViewController.swift
View file
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
29 changes: 18 additions & 11 deletions PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift
View file
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,8 +244,9 @@ 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
Expand All @@ -270,19 +271,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
2 changes: 1 addition & 1 deletion PoseEstimation-TFLiteSwift/StillImageHeatmapViewController.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ class StillImageHeatmapViewController: UIViewController {
@IBOutlet weak var pafSegment: UISegmentedControl?

// MARK: - ML Property
let poseEstimator: PoseEstimator = OpenPosePoseEstimator()
let poseEstimator: PoseEstimator = IMGCLSPoseEstimator()
var modelOutput: TFLiteFlatArray<Float32>? {
didSet {
updateHeatmapOverlayView()
Expand Down
2 changes: 1 addition & 1 deletion PoseEstimation-TFLiteSwift/StillImageLineViewController.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ class StillImageLineViewController: UIViewController {
}

// MARK: - ML Property
let poseEstimator: PoseEstimator = OpenPosePoseEstimator()
let poseEstimator: PoseEstimator = IMGCLSPoseEstimator()
var outputHumans: [PoseEstimationOutput.Human] = [] {
didSet {
updateOverlayView()
Expand Down
Loading

0 comments on commit 5dae993

Please sign in to comment.