[PR] Add some control UI components for parameters of post-process

PoseEstimation-TFLiteSwift.xcodeproj/project.pbxproj

@@ -30,6 +30,7 @@
 		7138DCCF242142FE0048E1D2 /* TFLiteFlatArray.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */; };
 		71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED1E24527D55001F796C /* PoseConfidenceMapDrawingView.swift */; };
 		71A1ED4124574F2E001F796C /* StillImageHeatmapViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71A1ED4024574F2E001F796C /* StillImageHeatmapViewController.swift */; };
+		71B07B97245E5C6C001FD385 /* NumericExtension.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71B07B96245E5C6C001FD385 /* NumericExtension.swift */; };
 		71DD577F2446D7CF0024C146 /* NonMaximumnonSuppression.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71DD577E2446D7CF0024C146 /* NonMaximumnonSuppression.swift */; };
 		71E8D9172438BAC10081DD6E /* openpose_ildoonet.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 71E8D9162438BAC10081DD6E /* openpose_ildoonet.tflite */; };
 		71E8D9192438BAD80081DD6E /* OpenPosePoseEstimator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 71E8D9182438BAD80081DD6E /* OpenPosePoseEstimator.swift */; };
@@ -65,6 +66,7 @@
 		7138DCCE242142FE0048E1D2 /* TFLiteFlatArray.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TFLiteFlatArray.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>"; };
+		71B07B96245E5C6C001FD385 /* NumericExtension.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NumericExtension.swift; sourceTree = "<group>"; };
 		71DD577E2446D7CF0024C146 /* NonMaximumnonSuppression.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NonMaximumnonSuppression.swift; sourceTree = "<group>"; };
 		71E8D9162438BAC10081DD6E /* openpose_ildoonet.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = openpose_ildoonet.tflite; sourceTree = "<group>"; };
 		71E8D9182438BAD80081DD6E /* OpenPosePoseEstimator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OpenPosePoseEstimator.swift; sourceTree = "<group>"; };
@@ -138,6 +140,7 @@
 				7105C93B241E8CE3001A4325 /* CVPixelBufferExtension.swift */,
 				712A7FC52425FD7200B043F9 /* UIImageExtension.swift */,
 				7105C93D241E90C2001A4325 /* DataExtension.swift */,
+				71B07B96245E5C6C001FD385 /* NumericExtension.swift */,
 			);
 			name = Extension;
 			sourceTree = "<group>";
@@ -325,6 +328,7 @@
 				71A1ED1F24527D55001F796C /* PoseConfidenceMapDrawingView.swift in Sources */,
 				7105C93C241E8CE3001A4325 /* CVPixelBufferExtension.swift in Sources */,
 				712A7FC9242667C900B043F9 /* PEFMHourglassPoseEstimator.swift in Sources */,
+				71B07B97245E5C6C001FD385 /* NumericExtension.swift in Sources */,
 				7105C91A241CE9B6001A4325 /* LiveImageViewController.swift in Sources */,
 				7105C916241CE9B5001A4325 /* AppDelegate.swift in Sources */,
 				7105C92F241D0235001A4325 /* PoseEstimator.swift in Sources */,

PoseEstimation-TFLiteSwift/LiveImageViewController.swift

@@ -14,12 +14,26 @@ class LiveImageViewController: UIViewController {
     // MARK: - IBOutlets
     @IBOutlet weak var previewView: UIView?
     @IBOutlet weak var overlayLineDotView: PoseKeypointsDrawingView?
+    @IBOutlet weak var humanTypeSegment: UISegmentedControl?
+    @IBOutlet weak var dimensionSegment: UISegmentedControl?
     @IBOutlet var partButtons: [UIButton]?
-    @IBOutlet weak var thresholdLabel: UILabel?
-    @IBOutlet weak var thresholdSlider: UISlider?
+    @IBOutlet weak var partThresholdLabel: UILabel?
+    @IBOutlet weak var partThresholdSlider: UISlider?
+    @IBOutlet weak var pairThresholdLabel: UILabel?
+    @IBOutlet weak var pairThresholdSlider: UISlider?
+    @IBOutlet weak var pairNMSFilterSizeLabel: UILabel?
+    @IBOutlet weak var pairNMSFilterSizeStepper: UIStepper?
+    @IBOutlet weak var humanMaxNumberLabel: UILabel?
+    @IBOutlet weak var humanMaxNumberStepper: UIStepper?
 
     var overlayViewRelativeRect: CGRect = .zero
+    var pixelBufferWidth: CGFloat = 0
 
+    var isSinglePerson: Bool = true {
+        didSet {
+            humanTypeSegment?.selectedSegmentIndex = isSinglePerson ? 0 : 1
+        }
+    }
     lazy var partIndexes: [String: Int] = {
         var partIndexes: [String: Int] = [:]
         poseEstimator.partNames.enumerated().forEach { offset, partName in
@@ -32,16 +46,58 @@ class LiveImageViewController: UIViewController {
         guard let partName = selectedPartName.components(separatedBy: "(").first else { return nil }
         return partIndexes[partName]
     }
-    var threshold: Float? {
+    var partThreshold: Float? {
         didSet {
-            guard let thresholdSlider = thresholdSlider else { return }
-            if let threshold = threshold {
-                thresholdSlider.value = threshold
-            } else {
-                thresholdSlider.value = thresholdSlider.minimumValue
+            let (slider, label, value) = (partThresholdSlider, partThresholdLabel, partThreshold)
+            if let slider = slider { slider.value = value ?? slider.minimumValue }
+            if let label = label { label.text = value.labelString }
+        }
+    }
+    var pairThreshold: Float? {
+        didSet {
+            let (slider, label, value) = (pairThresholdSlider, pairThresholdLabel, pairThreshold)
+            if let slider = slider { slider.value = value ?? slider.minimumValue }
+            if let label = label { label.text = value.labelString }
+        }
+    }
+    var pairNMSFilterSize: Int = 3 {
+        didSet {
+            let (stepper, label, value) = (pairNMSFilterSizeStepper, pairNMSFilterSizeLabel, pairNMSFilterSize)
+            if let stepper = stepper { stepper.value = Double(value) }
+            if let label = label { label.text = value.labelString }
+        }
+    }
+    var humanMaxNumber: Int? = 5 {
+        didSet {
+            let (stepper, label, value) = (humanMaxNumberStepper, humanMaxNumberLabel, humanMaxNumber)
+            if let stepper = stepper {
+                guard Int(stepper.minimumValue) != value else { humanMaxNumber = nil; return }
+                if let value = value { stepper.value = Double(value) }
+                else { stepper.value = stepper.minimumValue }
             }
+            if let label = label { label.text = value.labelString }
+        }
+    }
+
+    var preprocessOptions: PreprocessOptions {
+        let scalingRatio = pixelBufferWidth / overlayViewRelativeRect.width
+        let targetAreaRect = overlayViewRelativeRect.scaled(to: scalingRatio)
+        return PreprocessOptions(cropArea: .customAspectFill(rect: targetAreaRect))
+    }
+    var humanType: PostprocessOptions.HumanType {
+        if isSinglePerson {
+            return .singlePerson
+        } else {
+            return .multiPerson(pairThreshold: pairThreshold,
+                                nmsFilterSize: pairNMSFilterSize,
+                                maxHumanNumber: humanMaxNumber)
         }
     }
+    var postprocessOptions: PostprocessOptions {
+        return PostprocessOptions(partThreshold: partThreshold,
+                                  bodyPart: selectedPartIndex,
+                                  humanType: humanType)
+    }
 
     // MARK: - VideoCapture Properties
     var videoCapture = VideoCapture()
@@ -59,7 +115,12 @@ class LiveImageViewController: UIViewController {
         setUpUI()
 
         // setup initial post-process params
-        threshold = 0.1 // initial threshold for part (not for pair)
+        isSinglePerson = true   /// `multi-pose`
+        partThreshold = 0.1     /// 
+        pairThreshold = 3.4     /// Only used on `multi-person` mode. Before sort edges by cost, filter by pairThreshold for performance
+        pairNMSFilterSize = 3   /// Only used on `multi-person` mode. If 3, real could be 7X7 filter // (3●2+1)X(3●2+1)
+        humanMaxNumber = nil    /// Only used on `multi-person` mode. Not support yet
+
         select(on: "ALL")
     }
 
@@ -121,7 +182,7 @@ class LiveImageViewController: UIViewController {
             button.addTarget(self, action: #selector(selectPart), for: .touchUpInside)
         }
 
-        thresholdSlider?.isContinuous = false // `changeThreshold` will be called when touch up on slider
+        partThresholdSlider?.isContinuous = false // `changeThreshold` will be called when touch up on slider
     }
 
     override func viewDidLayoutSubviews() {
@@ -162,13 +223,28 @@ class LiveImageViewController: UIViewController {
         updatePartButton(on: partName)
     }
 
-    @IBAction func didChangedThresholdValue(_ sender: UISlider) {
-        threshold = (sender.value == sender.minimumValue) ? nil : sender.value
-        if let threshold = threshold {
-            thresholdLabel?.text = String(format: "%.2f", threshold)
-        } else {
-            thresholdLabel?.text = "nil"
-        }
+    @IBAction func didChangeHumanType(_ sender: UISegmentedControl) {
+        isSinglePerson = (sender.selectedSegmentIndex == 0)
+    }
+
+    @IBAction func didChangeDimension(_ sender: UISegmentedControl) {
+        // NOT SUPPORT YET
+    }
+
+    @IBAction func didChangedPartThreshold(_ sender: UISlider) {
+        partThreshold = (sender.value == sender.minimumValue) ? nil : sender.value
+    }
+
+    @IBAction func didChangePairThreshold(_ sender: UISlider) {
+        pairThreshold = (sender.value == sender.minimumValue) ? nil : sender.value
+    }
+
+    @IBAction func didChangePairNMSFilterSize(_ sender: UIStepper) {
+        pairNMSFilterSize = Int(sender.value)
+    }
+
+    @IBAction func didChangeHumanMaxNumber(_ sender: UIStepper) {
+        humanMaxNumber = (sender.value == sender.minimumValue) ? nil : Int(sender.value)
     }
 }
 
@@ -181,12 +257,11 @@ extension LiveImageViewController: VideoCaptureDelegate {
 
 extension LiveImageViewController {
     func inference(with pixelBuffer: CVPixelBuffer) {
-        let scalingRatio = pixelBuffer.size.width / overlayViewRelativeRect.width
-        let targetAreaRect = overlayViewRelativeRect.scaled(to: scalingRatio)
-        let partIndex: Int? = selectedPartIndex
-        let threshold: Float? = self.threshold
-        let input: PoseEstimationInput = .pixelBuffer(pixelBuffer: pixelBuffer, cropArea: .customAspectFill(rect: targetAreaRect))
-        let result: Result<PoseEstimationOutput, PoseEstimationError> = poseEstimator.inference(input, with: threshold, on: partIndex)
+        pixelBufferWidth = pixelBuffer.size.width
+        let input: PoseEstimationInput = .pixelBuffer(pixelBuffer: pixelBuffer,
+                                                      preprocessOptions: preprocessOptions,
+                                                      postprocessOptions: postprocessOptions)
+        let result: Result<PoseEstimationOutput, PoseEstimationError> = poseEstimator.inference(input)
 
         switch (result) {
         case .success(let output):

PoseEstimation-TFLiteSwift/NumericExtension.swift

@@ -0,0 +1,42 @@
+//
+//  NumericExtension.swift
+//  PoseEstimation-TFLiteSwift
+//
+//  Created by Doyoung Gwak on 2020/05/03.
+//  Copyright © 2020 Doyoung Gwak. All rights reserved.
+//
+
+import Foundation
+
+extension Optional where Wrapped == Float {
+    var labelString: String {
+        guard let value = self else { return "nil" }
+        return String(format: "%.2f", value)
+    }
+}
+
+extension Float32 {
+    func string(_ format: String = "%.2f") -> String {
+        return String(format: format, self)
+    }
+}
+
+extension Optional where Wrapped == Float {
+    static func *(lhs: Wrapped?, rhs: Float) -> Self {
+        guard let lhs = lhs else { return nil }
+        return some(lhs * rhs)
+    }
+}
+
+extension Int {
+    var labelString: String {
+        return String(format: "%d", self)
+    }
+}
+
+extension Optional where Wrapped == Int {
+    var labelString: String {
+        guard let value = self else { return "nil" }
+        return value.labelString
+    }
+}