feat(android): T-AND-003 MediaPipe Pose 端侧推理

2026-09-14 20:40:16 +08:00
parent 7a8dcfba25
commit 30c316a1cc
1 changed files with 219 additions and 0 deletions
@@ -0,0 +1,219 @@
+package com.opendance.studio.pose
+
+import android.content.Context
+import android.graphics.Bitmap
+import android.util.Log
+import androidx.camera.core.ImageProxy
+import com.google.mediapipe.framework.image.BitmapImageBuilder
+import com.google.mediapipe.framework.image.MPImage
+import com.google.mediapipe.tasks.core.BaseOptions
+import com.google.mediapipe.tasks.core.Delegate
+import com.google.mediapipe.tasks.vision.core.RunningMode
+import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarker
+import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarkerResult
+import com.opendance.studio.model.KeypointFrame
+import com.opendance.studio.model.SkeletonTopology
+
+/**
+ * T-AND-003：MediaPipe Pose 端侧推理
+ *
+ * 使用 LIVE_STREAM 模式：`detectAsync` 立即返回，结果通过监听器异步回调。
+ * 模型 `pose_landmarker_lite.task` 随 APK 分发（assets/，不压缩）。
+ *
+ * 输出 33 个归一化关键点，直接对齐契约 contracts/skeleton-topology.md §1
+ * 的索引顺序（MediaPipe Pose 的原始顺序即是契约顺序）。
+ */
+class PoseDetector(private val context: Context) {
+
+    /** 一帧推理结果 */
+    data class PoseResult(
+        val frame: KeypointFrame.PoseFrame,
+        /** 供叠加层绘制：归一化坐标 [x, y]，与预览同一坐标系（已按需镜像） */
+        val overlay: FloatArray,
+        val inferenceMs: Long,
+        val timestampMs: Long,
+    )
+
+    private var landmarker: PoseLandmarker? = null
+    private var lastTimestampMs = 0L
+    private var busy = false
+
+    var onResult: ((PoseResult) -> Unit)? = null
+    var onError: ((String) -> Unit)? = null
+
+    /** 是否把 x 翻转为镜像坐标（前置摄像头预览默认镜像，需与预览一致） */
+    @Volatile var mirror: Boolean = true
+
+    /** 仅上半身模式：下半身点置信度置 0，接收端会插值 */
+    @Volatile var upperBodyOnly: Boolean = false
+
+    val isReady: Boolean get() = landmarker != null
+
+    fun initialize(
+        modelAssetPath: String = DEFAULT_MODEL,
+        numPoses: Int = 1,
+        delegate: Delegate = Delegate.CPU,
+    ): Boolean {
+        if (landmarker != null) return true
+        return try {
+            val baseOptions = BaseOptions.builder()
+                .setModelAssetPath(modelAssetPath)
+                .setDelegate(delegate)
+                .build()
+
+            val options = PoseLandmarker.PoseLandmarkerOptions.builder()
+                .setBaseOptions(baseOptions)
+                .setRunningMode(RunningMode.LIVE_STREAM)
+                .setNumPoses(numPoses)
+                .setMinPoseDetectionConfidence(0.5f)
+                .setMinPosePresenceConfidence(0.5f)
+                .setMinTrackingConfidence(0.5f)
+                .setResultListener { result, input -> handleResult(result, input) }
+                .setErrorListener { error ->
+                    Log.e(TAG, "推理错误：${error.message}")
+                    onError?.invoke(error.message ?: "推理错误")
+                }
+                .build()
+
+            landmarker = PoseLandmarker.createFromOptions(context, options)
+            Log.i(TAG, "PoseLandmarker 就绪（$modelAssetPath，delegate=$delegate）")
+            true
+        } catch (e: Exception) {
+            Log.e(TAG, "PoseLandmarker 初始化失败", e)
+            onError?.invoke("MediaPipe 初始化失败：${e.message}")
+            landmarker = null
+            false
+        }
+    }
+
+    /**
+     * 提交一帧做异步推理；本方法不阻塞调用线程，且**不负责关闭 imageProxy**。
+     *
+     * @return 是否成功提交（上一帧仍在推理时返回 false，由调用方丢弃该帧）
+     */
+    fun detect(imageProxy: ImageProxy, timestampMs: Long, rotationDegrees: Int = 0): Boolean {
+        val detector = landmarker ?: return false
+        if (busy) return false
+
+        return try {
+            val bitmap = imageProxy.toBitmap()
+            val rotated = if (rotationDegrees != 0) rotateBitmap(bitmap, rotationDegrees) else bitmap
+            if (rotated !== bitmap) bitmap.recycle()
+
+            val mpImage: MPImage = BitmapImageBuilder(rotated).build()
+            val safeTimestamp = if (timestampMs <= lastTimestampMs) lastTimestampMs + 1 else timestampMs
+            busy = true
+            detector.detectAsync(mpImage, safeTimestamp)
+            true
+        } catch (e: Exception) {
+            busy = false
+            Log.w(TAG, "提交帧失败：${e.message}")
+            false
+        }
+    }
+
+    private fun handleResult(result: PoseLandmarkerResult, input: MPImage) {
+        busy = false
+        val started = System.currentTimeMillis()
+
+        val landmarks = result.landmarks().firstOrNull()
+        if (landmarks.isNullOrEmpty()) {
+            // 未检测到人体：返回空结果让上层显示「未识别到人体」
+            onResult?.invoke(
+                PoseResult(
+                    frame = emptyFrame(),
+                    overlay = FloatArray(0),
+                    inferenceMs = 0,
+                    timestampMs = lastTimestampMs,
+                ),
+            )
+            return
+        }
+
+        val count = minOf(landmarks.size, SkeletonTopology.KEYPOINT_COUNT)
+        val xs = FloatArray(SkeletonTopology.KEYPOINT_COUNT)
+        val ys = FloatArray(SkeletonTopology.KEYPOINT_COUNT)
+        val zs = FloatArray(SkeletonTopology.KEYPOINT_COUNT)
+        val conf = FloatArray(SkeletonTopology.KEYPOINT_COUNT)
+        val overlay = FloatArray(SkeletonTopology.KEYPOINT_COUNT * 2)
+
+        for (i in 0 until SkeletonTopology.KEYPOINT_COUNT) {
+            if (i >= count) {
+                // 缺失点：置信度 0，坐标归零，接收端会插值
+                xs[i] = 0.5f; ys[i] = 0.5f; zs[i] = 0f; conf[i] = 0f
+                overlay[i * 2] = 0.5f; overlay[i * 2 + 1] = 0.5f
+                continue
+            }
+            val lm = landmarks[i]
+            val rawX = lm.x().coerceIn(0f, 1f)
+            val y = lm.y().coerceIn(0f, 1f)
+            val z = lm.z()
+            // MediaPipe 的 visibility 可能为 null，缺失时按 presence 兜底
+            val visibility = lm.visibility().orElse(1f).coerceIn(0f, 1f)
+
+            val isLowerBody = i >= SkeletonTopology.LEFT_HIP && i != SkeletonTopology.NOSE
+            val effectiveConfidence = if (upperBodyOnly && isLowerBody) 0f else visibility
+
+            // 契约坐标：x,y ∈ [0,1] 图像空间；发送**未镜像**的原始坐标，
+            // 是否镜像由 flags.bit0 声明，接收端据此决定如何解释。
+            xs[i] = rawX
+            ys[i] = y
+            zs[i] = z
+            conf[i] = effectiveConfidence
+
+            // 叠加层坐标需要与预览一致：前置摄像头预览镜像，故叠加层 x 取反
+            val displayX = if (mirror) 1f - rawX else rawX
+            overlay[i * 2] = displayX
+            overlay[i * 2 + 1] = y
+        }
+
+        var flags = 0
+        if (mirror) flags = flags or KeypointFrame.FLAG_MIRRORED
+        if (upperBodyOnly) flags = flags or KeypointFrame.FLAG_UPPER_BODY_ONLY
+        flags = KeypointFrame.withLowConfidenceFlag(flags, conf)
+
+        val frame = KeypointFrame.PoseFrame(
+            frameId = (lastTimestampMs and 0x7FFFFFFF).toInt(),
+            flags = flags,
+            x = xs, y = ys, z = zs, confidence = conf,
+        )
+
+        onResult?.invoke(
+            PoseResult(
+                frame = frame,
+                overlay = overlay,
+                inferenceMs = System.currentTimeMillis() - started,
+                timestampMs = lastTimestampMs,
+            ),
+        )
+    }
+
+    private fun emptyFrame(): KeypointFrame.PoseFrame = KeypointFrame.PoseFrame(
+        frameId = 0,
+        flags = 0,
+        x = FloatArray(SkeletonTopology.KEYPOINT_COUNT) { 0.5f },
+        y = FloatArray(SkeletonTopology.KEYPOINT_COUNT) { 0.5f },
+        z = FloatArray(SkeletonTopology.KEYPOINT_COUNT),
+        confidence = FloatArray(SkeletonTopology.KEYPOINT_COUNT),
+    )
+
+    private fun rotateBitmap(source: Bitmap, degrees: Int): Bitmap {
+        val matrix = android.graphics.Matrix().apply { postRotate(degrees.toFloat()) }
+        return Bitmap.createBitmap(source, 0, 0, source.width, source.height, matrix, true)
+    }
+
+    fun close() {
+        try {
+            landmarker?.close()
+        } catch (e: Exception) {
+            Log.w(TAG, "关闭 PoseLandmarker 失败：${e.message}")
+        }
+        landmarker = null
+        busy = false
+    }
+
+    companion object {
+        private const val TAG = "PoseDetector"
+        const val DEFAULT_MODEL = "pose_landmarker_lite.task"
+    }
+}