Yolo pose estimation output processing in Android Tensorflow Lite API

Question

Im trying to implement yolo11n-pose_float16.tflite model to Android Kotlin project. But I cant get correct keypoints on output because there is something I did wrong. I want to draw only the skeleton containing keypoints and body part connections on the screen. How should I extract the output array to my Person object?

data class Person(
 var keyPoints: MutableList<KeyPoint>,
 val score: Float
)
enum class BodyPart {
 NOSE, LEFT_EYE, RIGHT_EYE, LEFT_EAR, RIGHT_EAR,
 LEFT_SHOULDER, RIGHT_SHOULDER, LEFT_ELBOW, RIGHT_ELBOW,
 LEFT_WRIST, RIGHT_WRIST, LEFT_HIP, RIGHT_HIP,
 LEFT_KNEE, RIGHT_KNEE, LEFT_ANKLE, RIGHT_ANKLE
}
data class KeyPoint(val bodyPart: BodyPart, var coordinate: PointF, val score: Float)
// Here is my function:
private fun processOutput(output: FloatArray): List<Person> {
 
 val persons = mutableListOf<Person>()
 // How should I process output???
 // output shape is [1,56,8400]
 
}

Answer 1

I figure it out, here post processing for 1,56,8400 shape output

for (c in 0 until numElements) {
 val cnf = array[c + numElements * 4]
 if (cnf > CONFIDENCE_THRESHOLD) {
 val cx = array[c]
 val cy = array[c + numElements]
 val w = array[c + numElements * 2]
 val h = array[c + numElements * 3]
 val x1 = cx - (w / 2F)
 val y1 = cy - (h / 2F)
 val x2 = cx + (w / 2F)
 val y2 = cy + (h / 2F)
 if (x1 <= 0F || x1 >= tensorWidth) continue
 if (y1 <= 0F || y1 >= tensorHeight) continue
 if (x2 <= 0F || x2 >= tensorWidth) continue
 if (y2 <= 0F || y2 >= tensorHeight) continue
 val keypoints = mutableListOf<KeyPoint>()
 for (k in 0 until 17) {
 var kx = array[c + numElements * (5 + k * 3)]
 var ky = array[c + numElements * (5 + k * 3 + 1)]
 kx /= tensorWidth
 ky /= tensorHeight
 keypoints.add(KeyPoint(kx, ky))
 }
 boundingBoxes.add(
 BoundingBox(
 x1 = x1, y1 = y1, x2 = x2, y2 = y2,
 cx = cx, cy = cy, w = w, h = h, cnf = cnf,
 keyPoints = keypoints
 )
 )
 }
 }

And here for body connections for this,

private val edges = listOf(
 Pair(0, 1), Pair(0, 2), // Nose ↔ Left Eye, Nose ↔ Right Eye
 Pair(1, 3), Pair(2, 4), // Left Eye ↔ Left Ear, Right Eye ↔ Right Ear
 Pair(0, 5), Pair(0, 6), // Nose ↔ Left Shoulder, Nose ↔ Right Shoulder
 Pair(5, 7), Pair(7, 9), // Left Shoulder ↔ Left Elbow, Left Elbow ↔ Left Wrist
 Pair(6, 8), Pair(8, 10), // Right Shoulder ↔ Right Elbow, Right Elbow ↔ Right Wrist
 Pair(5, 6), // Left Shoulder ↔ Right Shoulder
 Pair(5, 11), Pair(6, 12), // Left Shoulder ↔ Left Hip, Right Shoulder ↔ Right Hip
 Pair(11, 12), // Left Hip ↔ Right Hip
 Pair(11, 13), Pair(13, 15), // Left Hip ↔ Left Knee, Left Knee ↔ Left Ankle
 Pair(12, 14), Pair(14, 16) // Right Hip ↔ Right Knee, Right Knee ↔ Right Ankle
 )