Skip to content

Commit c31d8dd

Browse files
tucan9389tucan9389
authored
Merge pull request #66 from tucan9389/3d-pose
[PR] Implement post-process of 3d pose estimation (SimpleBaseline)
2 parents da2931e + 832c8be commit c31d8dd

21 files changed

Lines changed: 1280 additions & 324 deletions

‎.gitignore‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,4 +93,5 @@ Pods/*
9393
Podfile.lock
9494

9595
# TensorFlow Lite Model
96-
*.tflite
96+
*.tflite
97+
.DS_Store

‎Podfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,6 @@ target 'PoseEstimation-TFLiteSwift' do
66
use_frameworks!
77

88
# Pods for PoseEstimation-TFLiteSwift
9-
pod 'TensorFlowLiteSwift'
9+
pod 'TensorFlowLiteSwift/CoreML', '~> 2.4.0'
1010

1111
end

‎PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj‎

Lines changed: 69 additions & 43 deletions
Large diffs are not rendered by default.

‎PoseEstimation-TFLiteSwift/Base.lproj/Main.storyboard‎

Lines changed: 263 additions & 179 deletions
Large diffs are not rendered by default.

‎PoseEstimation-TFLiteSwift/Baseline3DPose/Baseline3DPoseEstimator.swift‎

Lines changed: 397 additions & 0 deletions
Large diffs are not rendered by default.

‎PoseEstimation-TFLiteSwift/CVPixelBufferExtension.swift‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ extension CVPixelBuffer {
115115
/// floating point values).
116116
/// - Returns: The RGB data representation of the image buffer or `nil` if the buffer could not be
117117
/// converted.
118-
func rgbData(byteCount: Int, isNormalized: Bool = false, isModelQuantized: Bool) -> Data? {
118+
func rgbData(byteCount: Int, normalization: TFLiteImageInterpreter.NormalizationOptions = .none, isModelQuantized: Bool) -> Data? {
119119
CVPixelBufferLockBaseAddress(self, .readOnly)
120120
defer { CVPixelBufferUnlockBaseAddress(self, .readOnly) }
121121
guard let sourceData = CVPixelBufferGetBaseAddress(self) else {
@@ -165,9 +165,19 @@ extension CVPixelBuffer {
165165
if isModelQuantized { return imageByteData }
166166

167167
let imageBytes = [UInt8](imageByteData)
168-
let bytes: [Float]
169-
if isNormalized {
168+
var bytes: [Float] = []
169+
if normalization == .scaledNormalization {
170170
bytes = imageBytes.map { Float($0) / 255.0 } // normalization
171+
} else if normalization == .pytorchNormalization {
172+
// bytes = imageBytes.map { Float($0) / 255.0 } // normalization
173+
bytes = imageBytes.map { Float($0) } // normalization
174+
for i in 0 ..< width * height {
175+
bytes[i ] = (Float32(imageBytes[i * 3 + 0]) - 0.485) / 0.229 // R
176+
bytes[width * height + i ] = (Float32(imageBytes[i * 3 + 1]) - 0.456) / 0.224 // G
177+
bytes[width * height * 2 + i] = (Float32(imageBytes[i * 3 + 2]) - 0.406) / 0.225 // B
178+
}
179+
} else if normalization == .meanStdNormalization {
180+
assert(false, "not support '.meanStdNormalization'")
171181
} else {
172182
bytes = imageBytes.map { Float($0) } // not normalization
173183
}

‎PoseEstimation-TFLiteSwift/LiveLineHeatmapViewController.swift‎

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

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

127127
override func viewDidLoad() {
128128
super.viewDidLoad()
@@ -292,10 +292,10 @@ extension LiveLineHeatmapViewController {
292292
DispatchQueue.main.async {
293293
if let partOffset = self.selectedPartIndex {
294294
self.lineDotView?.lines = []
295-
self.lineDotView?.keypoints = output.humans.map { $0.keypoints[partOffset] }
295+
self.lineDotView?.keypoints = output.humans2d.compactMap { $0 }.map { $0.keypoints[partOffset] }
296296
} else { // ALL case
297-
self.lineDotView?.lines = output.humans.reduce([]) { $0 + $1.lines }
298-
self.lineDotView?.keypoints = output.humans.reduce([]) { $0 + $1.keypoints }
297+
self.lineDotView?.lines = output.humans2d.compactMap { $0 }.reduce([]) { $0 + $1.lines }
298+
self.lineDotView?.keypoints = output.humans2d.compactMap { $0 }.reduce([]) { $0 + $1.keypoints }
299299
}
300300

301301
if let partOffset = self.selectedPartIndex {

‎PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ class OpenPosePoseEstimator: PoseEstimator {
3333
inputWidth: Input.width,
3434
inputHeight: Input.height,
3535
isGrayScale: Input.isGrayScale,
36-
isNormalized: Input.isNormalized
36+
normalization: Input.normalization
3737
)
3838
let imageInterpreter = TFLiteImageInterpreter(options: options)
3939
return imageInterpreter
@@ -88,7 +88,7 @@ private extension OpenPosePoseEstimator {
8888
static let width = 432
8989
static let height = 368
9090
static let isGrayScale = false
91-
static let isNormalized = false
91+
static let normalization = TFLiteImageInterpreter.NormalizationOptions.none
9292
}
9393
struct Output {
9494
struct ConfidenceMap { // similar to Heatmap
@@ -233,18 +233,18 @@ private extension PoseEstimationOutput {
233233
let human = parseSinglePerson(outputs,
234234
partIndex: postprocessOptions.bodyPart,
235235
partThreshold: postprocessOptions.partThreshold)
236-
humans = [human]
236+
humans = [.human2d(human: human)]
237237
case .multiPerson(let pairThreshold, let nmsFilterSize, let maxHumanNumber):
238238
humans = parseMultiHuman(outputs,
239239
partIndex: postprocessOptions.bodyPart,
240240
partThreshold: postprocessOptions.partThreshold,
241241
pairThreshold: pairThreshold,
242242
nmsFilterSize: nmsFilterSize,
243-
maxHumanNumber: maxHumanNumber)
243+
maxHumanNumber: maxHumanNumber).map { .human2d(human: $0) }
244244
}
245245
}
246246

247-
func parseSinglePerson(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?) -> Human {
247+
func parseSinglePerson(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?) -> Human2D {
248248
// openpose_ildoonet.tflite only use the first output
249249
let output = outputs[0]
250250

@@ -263,28 +263,28 @@ private extension PoseEstimationOutput {
263263
return (point: CGPoint(x: x, y: y), score: score)
264264
}
265265

266-
let keypoints: [Keypoint?] = keypointInfos
267-
.map { keypointInfo -> Keypoint? in Keypoint(position: keypointInfo.point, score: keypointInfo.score) }
268-
.map { keypointInfo -> Keypoint? in
266+
let keypoints: [Keypoint2D?] = keypointInfos
267+
.map { keypointInfo -> Keypoint2D? in Keypoint2D(position: keypointInfo.point, score: keypointInfo.score) }
268+
.map { keypointInfo -> Keypoint2D? in
269269
guard let score = keypointInfo?.score, let partThreshold = partThreshold else { return keypointInfo }
270270
return (score > partThreshold) ? keypointInfo : nil
271271
}
272272

273273
// lines
274-
var keypointWithBodyPart: [OpenPosePoseEstimator.Output.BodyPart: Keypoint] = [:]
274+
var keypointWithBodyPart: [OpenPosePoseEstimator.Output.BodyPart: Keypoint2D] = [:]
275275
OpenPosePoseEstimator.Output.BodyPart.allCases.enumerated().forEach { (index, bodyPart) in
276276
keypointWithBodyPart[bodyPart] = keypoints[index]
277277
}
278-
let lines: [Human.Line] = OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
278+
let lines: [Human2D.Line2D] = OpenPosePoseEstimator.Output.BodyPart.lines.compactMap { line in
279279
guard let fromKeypoint = keypointWithBodyPart[line.from],
280280
let toKeypoint = keypointWithBodyPart[line.to] else { return nil }
281281
return (from: fromKeypoint, to: toKeypoint)
282282
}
283283

284-
return Human(keypoints: keypoints, lines: lines)
284+
return Human2D(keypoints: keypoints, lines: lines)
285285
}
286286

287-
func parseMultiHuman(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?, pairThreshold: Float?, nmsFilterSize: Int, maxHumanNumber: Int?) -> [Human] {
287+
func parseMultiHuman(_ outputs: [TFLiteFlatArray<Float32>], partIndex: Int?, partThreshold: Float?, pairThreshold: Float?, nmsFilterSize: Int, maxHumanNumber: Int?) -> [Human2D] {
288288
// openpose_ildoonet.tflite only use the first output
289289
let output = outputs[0]
290290

@@ -303,15 +303,15 @@ private extension PoseEstimationOutput {
303303
}
304304
}
305305

306-
func parseSinglePartOnMultiHuman(_ output: TFLiteFlatArray<Float32>, partIndex: Int, partThreshold: Float?, nmsFilterSize: Int = 3) -> [Human] {
306+
func parseSinglePartOnMultiHuman(_ output: TFLiteFlatArray<Float32>, partIndex: Int, partThreshold: Float?, nmsFilterSize: Int = 3) -> [Human2D] {
307307
// process NMS
308308
let keypointIndexes = output.keypoints(partIndex: partIndex,
309309
filterSize: nmsFilterSize,
310310
threshold: partThreshold)
311311

312312
// convert col,row to Keypoint
313-
let kps: [Keypoint] = keypointIndexes.map { keypointInfo in
314-
return Keypoint(column: keypointInfo.col,
313+
let kps: [Keypoint2D] = keypointIndexes.map { keypointInfo in
314+
return Keypoint2D(column: keypointInfo.col,
315315
row: keypointInfo.row,
316316
width: OpenPosePoseEstimator.Output.ConfidenceMap.width,
317317
height: OpenPosePoseEstimator.Output.ConfidenceMap.height,
@@ -320,14 +320,14 @@ private extension PoseEstimationOutput {
320320

321321
// Make [Human]
322322
return kps.map { keypoint in
323-
let keypoints: [Keypoint?] = OpenPosePoseEstimator.Output.BodyPart.allCases.enumerated().map { offset, _ in
323+
let keypoints: [Keypoint2D?] = OpenPosePoseEstimator.Output.BodyPart.allCases.enumerated().map { offset, _ in
324324
return (offset == partIndex) ? keypoint : nil
325325
}
326-
return Human(keypoints: keypoints, lines: [])
326+
return Human2D(keypoints: keypoints, lines: [])
327327
}
328328
}
329329

330-
func parseAllPartOnMultiHuman(_ output: TFLiteFlatArray<Float32>, partIndex: Int?, partThreshold: Float?, pairThreshold: Float?, nmsFilterSize: Int, maxHumanNumber: Int?) -> [Human] {
330+
func parseAllPartOnMultiHuman(_ output: TFLiteFlatArray<Float32>, partIndex: Int?, partThreshold: Float?, pairThreshold: Float?, nmsFilterSize: Int, maxHumanNumber: Int?) -> [Human2D] {
331331

332332
let parts = OpenPosePoseEstimator.Output.BodyPart.allCases
333333
var verticesForEachPart: [[KeypointElement]?] = parts.map { _ in nil }
@@ -447,21 +447,21 @@ private extension PoseEstimationOutput {
447447
}
448448
}
449449

450-
let humans: [Human] = tmpHumans.map { tmpHuman in
451-
let keypoints: [Keypoint?] = tmpHuman.enumerated().map { (offset, locationInfo) in
450+
let humans: [Human2D] = tmpHumans.map { tmpHuman in
451+
let keypoints: [Keypoint2D?] = tmpHuman.enumerated().map { (offset, locationInfo) in
452452
guard let locationInfo = locationInfo else { return nil }
453-
return Keypoint(column: locationInfo.col,
453+
return Keypoint2D(column: locationInfo.col,
454454
row: locationInfo.row,
455455
width: colSize,
456456
height: rowSize,
457457
value: locationInfo.val)
458458
}
459-
let lines: [(from: Keypoint, to: Keypoint)] = pairs.compactMap { pair in
459+
let lines: [(from: Keypoint2D, to: Keypoint2D)] = pairs.compactMap { pair in
460460
guard let startingKeypoint = keypoints[pair.from.offsetValue()],
461461
let endingKeypoint = keypoints[pair.to.offsetValue()] else { return nil }
462462
return (from: startingKeypoint, to: endingKeypoint)
463463
}
464-
return Human(keypoints: keypoints, lines: lines)
464+
return Human2D(keypoints: keypoints, lines: lines)
465465
}
466466
return humans
467467
}
@@ -499,7 +499,7 @@ private extension TFLiteFlatArray where Element == Float32 {
499499
}
500500
}
501501

502-
private extension Keypoint {
502+
private extension Keypoint2D {
503503
init(column: Int, row: Int, width: Int, height: Int, value: Float32) {
504504
let x = (CGFloat(column) + 0.5) / CGFloat(width)
505505
let y = (CGFloat(row) + 0.5) / CGFloat(height)

0 commit comments

Comments
 (0)