173 lines
6.6 KiB
Kotlin
173 lines
6.6 KiB
Kotlin
package com.opendance.studio.model
|
||
|
||
import java.nio.ByteBuffer
|
||
import java.nio.ByteOrder
|
||
|
||
/**
|
||
* 239 字节关键点帧编解码 —— 契约 contracts/websocket-keypoints.md §2（冻结）
|
||
*
|
||
* 帧布局（**小端序**）：
|
||
* ```
|
||
* 偏移 长度 字段
|
||
* 0 2 magic u16 = 0x4F44 ('OD')
|
||
* 2 1 version u8 = 1
|
||
* 3 1 flags u8 bit0 镜像 | bit1 含低置信度点 | bit2 上半身模式
|
||
* 4 4 frame_id u32 LE，从 0 起单调递增
|
||
* 8 231 keypoints 33 × 7B：s16 x, s16 y, s16 z, u8 conf
|
||
* 合计 239
|
||
* ```
|
||
*
|
||
* 坐标语义（contracts/skeleton-topology.md §4）：
|
||
* x,y ∈ [0,1] 归一化图像坐标（原点左上，y 向下）；z 与 x 同尺度，
|
||
* 负值朝向相机。归一化系数统一 **32767**。
|
||
*
|
||
* 本实现与 opendance-web 的 TypeScript 参考实现、opendance-windows 的 GDScript
|
||
* 实现逐位对齐，并由单元测试用同一份金标准字节向量交叉校验。
|
||
*/
|
||
object KeypointFrame {
|
||
|
||
const val MAGIC = 0x4F44
|
||
const val PROTOCOL_VERSION = 1
|
||
const val BINARY_SIZE = 239
|
||
const val KEYPOINT_COUNT = 33
|
||
const val HEADER_SIZE = 8
|
||
const val POINT_BYTES = 7
|
||
const val SCALE = 32767.0
|
||
|
||
const val FLAG_MIRRORED = 1 shl 0
|
||
const val FLAG_LOW_CONFIDENCE = 1 shl 1
|
||
const val FLAG_UPPER_BODY_ONLY = 1 shl 2
|
||
|
||
/** 低于该置信度的点坐标无意义，接收端必须插值 */
|
||
const val CONFIDENCE_THRESHOLD = 0.5f
|
||
|
||
/** 一帧姿态数据（并行数组，避免 30 FPS 循环里的对象分配） */
|
||
class PoseFrame(
|
||
val frameId: Int,
|
||
val flags: Int,
|
||
val x: FloatArray,
|
||
val y: FloatArray,
|
||
val z: FloatArray,
|
||
val confidence: FloatArray,
|
||
) {
|
||
init {
|
||
require(x.size == KEYPOINT_COUNT) { "x 必须为 $KEYPOINT_COUNT 项" }
|
||
require(y.size == KEYPOINT_COUNT) { "y 必须为 $KEYPOINT_COUNT 项" }
|
||
require(z.size == KEYPOINT_COUNT) { "z 必须为 $KEYPOINT_COUNT 项" }
|
||
require(confidence.size == KEYPOINT_COUNT) { "confidence 必须为 $KEYPOINT_COUNT 项" }
|
||
}
|
||
|
||
val isMirrored: Boolean get() = (flags and FLAG_MIRRORED) != 0
|
||
val isUpperBodyOnly: Boolean get() = (flags and FLAG_UPPER_BODY_ONLY) != 0
|
||
val hasLowConfidence: Boolean get() = (flags and FLAG_LOW_CONFIDENCE) != 0
|
||
|
||
/** 该点是否可用于评分 */
|
||
fun isReliable(index: Int): Boolean = confidence[index] >= CONFIDENCE_THRESHOLD
|
||
|
||
/** 平均置信度，用于运行页质量指示 */
|
||
fun averageConfidence(): Float {
|
||
var total = 0f
|
||
for (c in confidence) total += c
|
||
return total / confidence.size
|
||
}
|
||
|
||
/** 已镜像的帧需翻转 x 才能与对端坐标系一致 */
|
||
fun xUnmirrored(index: Int): Float = if (isMirrored) 1f - x[index] else x[index]
|
||
}
|
||
|
||
/**
|
||
* 校验帧字节流；合法返回 null，否则返回可读原因。
|
||
*/
|
||
fun validate(bytes: ByteArray): String? {
|
||
if (bytes.size != BINARY_SIZE) {
|
||
return "帧长度必须为 $BINARY_SIZE 字节，实际 ${bytes.size}"
|
||
}
|
||
val buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
|
||
val magic = buffer.getShort(0).toInt() and 0xFFFF
|
||
if (magic != MAGIC) {
|
||
return "magic 校验失败：期望 0x%04X，实际 0x%04X".format(MAGIC, magic)
|
||
}
|
||
val version = buffer.get(2).toInt() and 0xFF
|
||
if (version != PROTOCOL_VERSION) {
|
||
return "不支持的协议版本 $version（本端支持 $PROTOCOL_VERSION）"
|
||
}
|
||
return null
|
||
}
|
||
|
||
/**
|
||
* 编码为 239 字节帧 —— 本端（Android）的主要职责。
|
||
*
|
||
* @param xs 归一化 x，[0,1]
|
||
* @param ys 归一化 y，[0,1]
|
||
* @param zs 深度，约 [-1,1]，负值朝相机
|
||
* @param confidence [0,1]
|
||
* @param flags 见 FLAG_* 常量
|
||
*/
|
||
fun encode(
|
||
frameId: Int,
|
||
xs: FloatArray,
|
||
ys: FloatArray,
|
||
zs: FloatArray,
|
||
confidence: FloatArray,
|
||
flags: Int = 0,
|
||
): ByteArray {
|
||
require(xs.size == KEYPOINT_COUNT) { "xs 必须为 $KEYPOINT_COUNT 项" }
|
||
require(ys.size == KEYPOINT_COUNT) { "ys 必须为 $KEYPOINT_COUNT 项" }
|
||
require(zs.size == KEYPOINT_COUNT) { "zs 必须为 $KEYPOINT_COUNT 项" }
|
||
require(confidence.size == KEYPOINT_COUNT) { "confidence 必须为 $KEYPOINT_COUNT 项" }
|
||
|
||
val buffer = ByteBuffer.allocate(BINARY_SIZE).order(ByteOrder.LITTLE_ENDIAN)
|
||
buffer.putShort(0, MAGIC.toShort())
|
||
buffer.put(2, PROTOCOL_VERSION.toByte())
|
||
buffer.put(3, (flags and 0xFF).toByte())
|
||
buffer.putInt(4, frameId)
|
||
|
||
for (i in 0 until KEYPOINT_COUNT) {
|
||
val offset = HEADER_SIZE + i * POINT_BYTES
|
||
buffer.putShort(offset, toFixed16(xs[i]))
|
||
buffer.putShort(offset + 2, toFixed16(ys[i]))
|
||
buffer.putShort(offset + 4, toFixed16(zs[i]))
|
||
buffer.put(offset + 6, toConfidenceByte(confidence[i]))
|
||
}
|
||
return buffer.array()
|
||
}
|
||
|
||
/** 解码 239 字节帧；非法输入抛 IllegalArgumentException */
|
||
fun decode(bytes: ByteArray): PoseFrame {
|
||
validate(bytes)?.let { throw IllegalArgumentException(it) }
|
||
|
||
val buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
|
||
val frameId = buffer.getInt(4)
|
||
val flags = buffer.get(3).toInt() and 0xFF
|
||
|
||
val xs = FloatArray(KEYPOINT_COUNT)
|
||
val ys = FloatArray(KEYPOINT_COUNT)
|
||
val zs = FloatArray(KEYPOINT_COUNT)
|
||
val confidence = FloatArray(KEYPOINT_COUNT)
|
||
|
||
for (i in 0 until KEYPOINT_COUNT) {
|
||
val offset = HEADER_SIZE + i * POINT_BYTES
|
||
xs[i] = buffer.getShort(offset) / SCALE.toFloat()
|
||
ys[i] = buffer.getShort(offset + 2) / SCALE.toFloat()
|
||
zs[i] = buffer.getShort(offset + 4) / SCALE.toFloat()
|
||
confidence[i] = (buffer.get(offset + 6).toInt() and 0xFF) / 255f
|
||
}
|
||
return PoseFrame(frameId, flags, xs, ys, zs, confidence)
|
||
}
|
||
|
||
/** 依据置信度数组计算是否应置 FLAG_LOW_CONFIDENCE */
|
||
fun computeLowConfidence(confidence: FloatArray): Boolean =
|
||
confidence.any { it < CONFIDENCE_THRESHOLD }
|
||
|
||
fun withLowConfidenceFlag(flags: Int, confidence: FloatArray): Int =
|
||
if (computeLowConfidence(confidence)) flags or FLAG_LOW_CONFIDENCE else flags and FLAG_LOW_CONFIDENCE.inv()
|
||
|
||
private fun toFixed16(value: Float): Short {
|
||
val scaled = Math.round(value * SCALE)
|
||
return scaled.coerceIn(-32768L, 32767L).toShort()
|
||
}
|
||
|
||
private fun toConfidenceByte(value: Float): Byte =
|
||
Math.round(value.coerceIn(0f, 1f) * 255f).coerceIn(0, 255).toByte()
|
||
}
|