Merge pull request #42 from tucan9389/multi-pose

[PR] Implement building bipartite graph, calculate cost values of each edge, assign and make humans
tucan9389 · May 1, 2020 · 98d29ba · 98d29ba
2 parents d84aa66 + 32521ae
commit 98d29ba
Show file tree

Hide file tree

Showing 7 changed files with 341 additions and 78 deletions.
diff --git a/PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj b/PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj
@@ -111,10 +111,10 @@
 				712A7FC02424BDD800B043F9 /* StillImageLineViewController.swift */,
 				71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */,
 				7105C919241CE9B6001A4325 /* LiveImageViewController.swift */,
+				7105C929241D011F001A4325 /* View */,
 				7105C92B241D0150001A4325 /* MLModel */,
 				71DD577D2446D7A40024C146 /* Algorithm */,
 				7105C92A241D0144001A4325 /* Extension */,
-				7105C929241D011F001A4325 /* View */,
 				7105C938241D29C5001A4325 /* Video */,
 				7105C91E241CE9B7001A4325 /* Assets.xcassets */,
 				7105C920241CE9B7001A4325 /* LaunchScreen.storyboard */,

diff --git a/PoseEstimation-TFLiteSwift/NonMaximumnonSuppression.swift b/PoseEstimation-TFLiteSwift/NonMaximumnonSuppression.swift
@@ -9,12 +9,13 @@
 import Foundation
 
 class NonMaximumnonSuppression {
-    typealias MaximumPoint = (row: Int, col: Int, val: Float32)
+    typealias MaximumPoint = (col: Int, row: Int, val: Float32)
 
     static func process(_ heatmap: TFLiteFlatArray<Float32>, partIndex: Int, width: Int, height: Int) -> [MaximumPoint] {
         let filterSize = 3
         var lastMaximumColumns: [MaximumPoint?] = (0..<width).map { _ in nil }
         var results: [MaximumPoint] = []
+        results.reserveCapacity(20)
 
         for row in (0..<height) {
             for col in (0..<width) {
@@ -48,8 +49,8 @@ class NonMaximumnonSuppression {
                                 lastMaximumColumns[targetColumn] = nil
                             }
                         }
-                        results.append((row: lastMaximumPoint.row,
-                                        col: lastMaximumPoint.col,
+                        results.append((col: lastMaximumPoint.col,
+                                        row: lastMaximumPoint.row,
                                         val: lastMaximumPoint.val))
                     }
                 }
@@ -66,8 +67,8 @@ class NonMaximumnonSuppression {
                     lastMaximumColumns[targetColumn] = nil
                 }
             }
-            results.append((row: lastMaximumPoint.row,
-                            col: lastMaximumPoint.col,
+            results.append((col: lastMaximumPoint.col,
+                            row: lastMaximumPoint.row,
                             val: lastMaximumPoint.val))
         }
 

diff --git a/PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift b/PoseEstimation-TFLiteSwift/OpenPose/OpenPosePoseEstimator.swift
diff --git a/PoseEstimation-TFLiteSwift/PoseConfidenceMapDrawingView.swift b/PoseEstimation-TFLiteSwift/PoseConfidenceMapDrawingView.swift
@@ -30,15 +30,16 @@ class PoseConfidenceMapDrawingView: UIView {
                 var componentOfVector = outputChannelIndexes.reduce(0.0) { value, outputChannelIndex in
                     return value + output[0, row, col, outputChannelIndex]
                 }
-                componentOfVector = min(max(componentOfVector, -1.0), 1.0)
+                componentOfVector = min(max(componentOfVector, -1.0), 1.0) // -1.0 ~ 1.0
                 let drawingAreaRect = CGRect(x: oneAreaWidth*CGFloat(col), y: oneAreaHeight*CGFloat(row),
                                              width: oneAreaWidth, height: oneAreaHeight)
-                drawRect(with: drawingAreaRect, componentOfVector: componentOfVector)
+                let areaFillColor = DrawingConstant.Area.areaColor(CGFloat(componentOfVector))
+                drawRect(with: drawingAreaRect, fillColor: areaFillColor)
             }
         }
     }
 
-    func drawRect(with rect: CGRect, componentOfVector: Float32 = 0.0) {
+    func drawRect(with rect: CGRect, fillColor: UIColor) {
         guard let startingPoint = rect.points.first else { return }
         let rectPath = UIBezierPath()
         rectPath.move(to: startingPoint)
@@ -51,7 +52,7 @@ class PoseConfidenceMapDrawingView: UIView {
         rectPath.lineWidth = DrawingConstant.Line.width
         DrawingConstant.Area.lineColor.setStroke()
         rectPath.stroke()
-        DrawingConstant.Area.areaColor(CGFloat(componentOfVector)).setFill()
+        fillColor.setFill()
         rectPath.fill()
     }
 }

diff --git a/PoseEstimation-TFLiteSwift/PoseEstimator.swift b/PoseEstimation-TFLiteSwift/PoseEstimator.swift
@@ -72,6 +72,22 @@ struct Keypoint {
     let score: Float
 }
 
+struct KeypointElement: Equatable {
+    let col: Int
+    let row: Int
+    let val: Float32
+
+    init(element: (col: Int, row: Int, val: Float32)) {
+        col = element.col
+        row = element.row
+        val = element.val
+    }
+
+    static func == (lhs: KeypointElement, rhs: KeypointElement) -> Bool {
+        return lhs.col == rhs.col && lhs.row == rhs.row
+    }
+}
+
 struct PoseEstimationOutput {
 
     var outputs: [TFLiteFlatArray<Float32>]

diff --git a/PoseEstimation-TFLiteSwift/StillImageHeatmapViewController.swift b/PoseEstimation-TFLiteSwift/StillImageHeatmapViewController.swift
@@ -137,7 +137,12 @@ class StillImageHeatmapViewController: UIViewController {
     }
 
     @IBAction func didChangePAF(_ sender: Any) {
-        updateHeatmapOverlayView()
+        if case .paf = selectedChannel {
+            updateHeatmapOverlayView()
+        } else {
+            guard let allButton = pairButtons?.first else { return }
+            selectChannel(allButton)
+        }
     }
 
     @objc func selectChannel(_ button: UIButton) {
@@ -148,11 +153,13 @@ class StillImageHeatmapViewController: UIViewController {
             updateChannelButton(with: partName, on: partButtons)
             updateChannelButton(with: nil, on: pairButtons)
             selectedChannel = .confidenceMap(part: partName)
+            UIView.animate(withDuration: 0.23) { self.pafSegment?.alpha = 0.5 }
         } else if pairButtons?.contains(button) == true {
             let pairName = buttonTitle.components(separatedBy: "(").first ?? buttonTitle
             updateChannelButton(with: nil, on: partButtons)
             updateChannelButton(with: pairName, on: pairButtons)
             selectedChannel = .paf(pair: pairName)
+            UIView.animate(withDuration: 0.23) { self.pafSegment?.alpha = 1.0 }
         }
 
         updateHeatmapOverlayView()

diff --git a/PoseEstimation-TFLiteSwift/StillImageLineViewController.swift b/PoseEstimation-TFLiteSwift/StillImageLineViewController.swift
@@ -91,7 +91,8 @@ class StillImageLineViewController: UIViewController {
             button.addTarget(self, action: #selector(selectPart), for: .touchUpInside)
         }
 
-        threshold = 0.1
+        thresholdSlider?.isContinuous = false // `changeThreshold` will be called when touch up on slider
+        threshold = 0.1 // initial threshold for part (not for pair)
     }
 
     func updatePartButton(on targetPartName: String) {
@@ -109,12 +110,6 @@ class StillImageLineViewController: UIViewController {
 
     @objc func selectPart(_ button: UIButton) {
         guard let partName = button.title(for: .normal) else { return }
-//        guard partName != "ALL" else {
-//            let alert = UIAlertController(title: "Error", message: "Not support 'ALL' case on multi pose estimation", preferredStyle: .alert)
-//            alert.addAction(UIAlertAction(title: "Ok", style: .cancel))
-//            present(alert, animated: true)
-//            return
-//        }
 
         select(on: partName)
         if let image = imageView?.image {