|
|
|
|
@ -33,40 +33,41 @@ import kotlin.math.abs
|
|
|
|
|
import kotlin.math.max
|
|
|
|
|
import kotlin.math.min
|
|
|
|
|
|
|
|
|
|
// 定义模型类型，支持Lightning和Thunder两种
|
|
|
|
|
enum class ModelType {
|
|
|
|
|
Lightning,
|
|
|
|
|
Thunder
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// MoveNet类实现了PoseDetector接口，用于姿势估计
|
|
|
|
|
class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: GpuDelegate?) :
|
|
|
|
|
PoseDetector {
|
|
|
|
|
|
|
|
|
|
companion object {
|
|
|
|
|
private const val MIN_CROP_KEYPOINT_SCORE = .2f
|
|
|
|
|
private const val CPU_NUM_THREADS = 4
|
|
|
|
|
private const val MIN_CROP_KEYPOINT_SCORE = .2f // 关键点最低得分
|
|
|
|
|
private const val CPU_NUM_THREADS = 4 // CPU线程数
|
|
|
|
|
|
|
|
|
|
// Parameters that control how large crop region should be expanded from previous frames'
|
|
|
|
|
// body keypoints.
|
|
|
|
|
private const val TORSO_EXPANSION_RATIO = 1.9f
|
|
|
|
|
private const val BODY_EXPANSION_RATIO = 1.2f
|
|
|
|
|
// 控制如何从上一帧的人体关键点扩展裁剪区域的参数
|
|
|
|
|
private const val TORSO_EXPANSION_RATIO = 1.9f // 胸部扩展比率
|
|
|
|
|
private const val BODY_EXPANSION_RATIO = 1.2f // 身体扩展比率
|
|
|
|
|
|
|
|
|
|
// TFLite file names.
|
|
|
|
|
// TFLite模型文件名称
|
|
|
|
|
private const val LIGHTNING_FILENAME = "movenet_lightning.tflite"
|
|
|
|
|
private const val THUNDER_FILENAME = "movenet_thunder.tflite"
|
|
|
|
|
|
|
|
|
|
// allow specifying model type.
|
|
|
|
|
// 根据设备类型创建MoveNet实例
|
|
|
|
|
fun create(context: Context, device: Device, modelType: ModelType): MoveNet {
|
|
|
|
|
val options = Interpreter.Options()
|
|
|
|
|
var gpuDelegate: GpuDelegate? = null
|
|
|
|
|
options.setNumThreads(CPU_NUM_THREADS)
|
|
|
|
|
options.setNumThreads(CPU_NUM_THREADS) // 设置CPU线程数
|
|
|
|
|
when (device) {
|
|
|
|
|
Device.CPU -> {
|
|
|
|
|
}
|
|
|
|
|
Device.GPU -> {
|
|
|
|
|
gpuDelegate = GpuDelegate()
|
|
|
|
|
gpuDelegate = GpuDelegate() // 使用GPU加速
|
|
|
|
|
options.addDelegate(gpuDelegate)
|
|
|
|
|
}
|
|
|
|
|
Device.NNAPI -> options.setUseNNAPI(true)
|
|
|
|
|
Device.NNAPI -> options.setUseNNAPI(true) // 使用NNAPI加速
|
|
|
|
|
}
|
|
|
|
|
return MoveNet(
|
|
|
|
|
Interpreter(
|
|
|
|
|
@ -80,26 +81,27 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// default to lightning.
|
|
|
|
|
// 默认使用Lightning模型
|
|
|
|
|
fun create(context: Context, device: Device): MoveNet =
|
|
|
|
|
create(context, device, ModelType.Lightning)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
private var cropRegion: RectF? = null
|
|
|
|
|
private var lastInferenceTimeNanos: Long = -1
|
|
|
|
|
private val inputWidth = interpreter.getInputTensor(0).shape()[1]
|
|
|
|
|
private val inputHeight = interpreter.getInputTensor(0).shape()[2]
|
|
|
|
|
private var outputShape: IntArray = interpreter.getOutputTensor(0).shape()
|
|
|
|
|
private var cropRegion: RectF? = null // 用于存储裁剪区域
|
|
|
|
|
private var lastInferenceTimeNanos: Long = -1 // 上次推理的时间
|
|
|
|
|
private val inputWidth = interpreter.getInputTensor(0).shape()[1] // 输入图像的宽度
|
|
|
|
|
private val inputHeight = interpreter.getInputTensor(0).shape()[2] // 输入图像的高度
|
|
|
|
|
private var outputShape: IntArray = interpreter.getOutputTensor(0).shape() // 输出的形状
|
|
|
|
|
|
|
|
|
|
// 估算姿势
|
|
|
|
|
override fun estimatePoses(bitmap: Bitmap): List<Person> {
|
|
|
|
|
val inferenceStartTimeNanos = SystemClock.elapsedRealtimeNanos()
|
|
|
|
|
val inferenceStartTimeNanos = SystemClock.elapsedRealtimeNanos() // 获取当前时间
|
|
|
|
|
if (cropRegion == null) {
|
|
|
|
|
cropRegion = initRectF(bitmap.width, bitmap.height)
|
|
|
|
|
cropRegion = initRectF(bitmap.width, bitmap.height) // 初始化裁剪区域
|
|
|
|
|
}
|
|
|
|
|
var totalScore = 0f
|
|
|
|
|
var totalScore = 0f // 总得分
|
|
|
|
|
|
|
|
|
|
val numKeyPoints = outputShape[2]
|
|
|
|
|
val keyPoints = mutableListOf<KeyPoint>()
|
|
|
|
|
val numKeyPoints = outputShape[2] // 关键点的数量
|
|
|
|
|
val keyPoints = mutableListOf<KeyPoint>() // 存储关键点
|
|
|
|
|
|
|
|
|
|
cropRegion?.run {
|
|
|
|
|
val rect = RectF(
|
|
|
|
|
@ -107,35 +109,35 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
(top * bitmap.height),
|
|
|
|
|
(right * bitmap.width),
|
|
|
|
|
(bottom * bitmap.height)
|
|
|
|
|
)
|
|
|
|
|
) // 根据裁剪区域计算裁剪矩形
|
|
|
|
|
val detectBitmap = Bitmap.createBitmap(
|
|
|
|
|
rect.width().toInt(),
|
|
|
|
|
rect.height().toInt(),
|
|
|
|
|
Bitmap.Config.ARGB_8888
|
|
|
|
|
)
|
|
|
|
|
) // 创建裁剪后的Bitmap
|
|
|
|
|
Canvas(detectBitmap).drawBitmap(
|
|
|
|
|
bitmap,
|
|
|
|
|
-rect.left,
|
|
|
|
|
-rect.top,
|
|
|
|
|
null
|
|
|
|
|
)
|
|
|
|
|
val inputTensor = processInputImage(detectBitmap, inputWidth, inputHeight)
|
|
|
|
|
val outputTensor = TensorBuffer.createFixedSize(outputShape, DataType.FLOAT32)
|
|
|
|
|
val inputTensor = processInputImage(detectBitmap, inputWidth, inputHeight) // 处理输入图像
|
|
|
|
|
val outputTensor = TensorBuffer.createFixedSize(outputShape, DataType.FLOAT32) // 创建输出张量
|
|
|
|
|
val widthRatio = detectBitmap.width.toFloat() / inputWidth
|
|
|
|
|
val heightRatio = detectBitmap.height.toFloat() / inputHeight
|
|
|
|
|
|
|
|
|
|
val positions = mutableListOf<Float>()
|
|
|
|
|
|
|
|
|
|
inputTensor?.let { input ->
|
|
|
|
|
interpreter.run(input.buffer, outputTensor.buffer.rewind())
|
|
|
|
|
interpreter.run(input.buffer, outputTensor.buffer.rewind()) // 运行推理
|
|
|
|
|
val output = outputTensor.floatArray
|
|
|
|
|
for (idx in 0 until numKeyPoints) {
|
|
|
|
|
val x = output[idx * 3 + 1] * inputWidth * widthRatio
|
|
|
|
|
val y = output[idx * 3 + 0] * inputHeight * heightRatio
|
|
|
|
|
val x = output[idx * 3 + 1] * inputWidth * widthRatio // 获取x坐标
|
|
|
|
|
val y = output[idx * 3 + 0] * inputHeight * heightRatio // 获取y坐标
|
|
|
|
|
|
|
|
|
|
positions.add(x)
|
|
|
|
|
positions.add(y)
|
|
|
|
|
val score = output[idx * 3 + 2]
|
|
|
|
|
val score = output[idx * 3 + 2] // 获取得分
|
|
|
|
|
keyPoints.add(
|
|
|
|
|
KeyPoint(
|
|
|
|
|
BodyPart.fromInt(idx),
|
|
|
|
|
@ -152,8 +154,8 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
val matrix = Matrix()
|
|
|
|
|
val points = positions.toFloatArray()
|
|
|
|
|
|
|
|
|
|
matrix.postTranslate(rect.left, rect.top)
|
|
|
|
|
matrix.mapPoints(points)
|
|
|
|
|
matrix.postTranslate(rect.left, rect.top) // 计算偏移
|
|
|
|
|
matrix.mapPoints(points) // 映射坐标
|
|
|
|
|
keyPoints.forEachIndexed { index, keyPoint ->
|
|
|
|
|
keyPoint.coordinate =
|
|
|
|
|
PointF(
|
|
|
|
|
@ -161,16 +163,18 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
points[index * 2 + 1]
|
|
|
|
|
)
|
|
|
|
|
}
|
|
|
|
|
// new crop region
|
|
|
|
|
// 更新裁剪区域
|
|
|
|
|
cropRegion = determineRectF(keyPoints, bitmap.width, bitmap.height)
|
|
|
|
|
}
|
|
|
|
|
lastInferenceTimeNanos =
|
|
|
|
|
SystemClock.elapsedRealtimeNanos() - inferenceStartTimeNanos
|
|
|
|
|
return listOf(Person(keyPoints = keyPoints, score = totalScore / numKeyPoints))
|
|
|
|
|
SystemClock.elapsedRealtimeNanos() - inferenceStartTimeNanos // 计算推理时间
|
|
|
|
|
return listOf(Person(keyPoints = keyPoints, score = totalScore / numKeyPoints)) // 返回姿势信息
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 获取上次推理时间
|
|
|
|
|
override fun lastInferenceTimeNanos(): Long = lastInferenceTimeNanos
|
|
|
|
|
|
|
|
|
|
// 关闭资源
|
|
|
|
|
override fun close() {
|
|
|
|
|
gpuDelegate?.close()
|
|
|
|
|
interpreter.close()
|
|
|
|
|
@ -178,27 +182,25 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Prepare input image for detection
|
|
|
|
|
* 准备输入图像进行检测
|
|
|
|
|
*/
|
|
|
|
|
private fun processInputImage(bitmap: Bitmap, inputWidth: Int, inputHeight: Int): TensorImage? {
|
|
|
|
|
val width: Int = bitmap.width
|
|
|
|
|
val height: Int = bitmap.height
|
|
|
|
|
|
|
|
|
|
val size = if (height > width) width else height
|
|
|
|
|
val size = if (height > width) width else height // 选择较小的一边
|
|
|
|
|
val imageProcessor = ImageProcessor.Builder().apply {
|
|
|
|
|
add(ResizeWithCropOrPadOp(size, size))
|
|
|
|
|
add(ResizeOp(inputWidth, inputHeight, ResizeOp.ResizeMethod.BILINEAR))
|
|
|
|
|
add(ResizeWithCropOrPadOp(size, size)) // 裁剪或填充为正方形
|
|
|
|
|
add(ResizeOp(inputWidth, inputHeight, ResizeOp.ResizeMethod.BILINEAR)) // 调整大小
|
|
|
|
|
}.build()
|
|
|
|
|
val tensorImage = TensorImage(DataType.UINT8)
|
|
|
|
|
tensorImage.load(bitmap)
|
|
|
|
|
return imageProcessor.process(tensorImage)
|
|
|
|
|
tensorImage.load(bitmap) // 加载图像
|
|
|
|
|
return imageProcessor.process(tensorImage) // 返回处理后的图像
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Defines the default crop region.
|
|
|
|
|
* The function provides the initial crop region (pads the full image from both
|
|
|
|
|
* sides to make it a square image) when the algorithm cannot reliably determine
|
|
|
|
|
* the crop region from the previous frame.
|
|
|
|
|
* 初始化裁剪区域。
|
|
|
|
|
* 当算法无法从上一帧可靠地确定裁剪区域时，提供初始裁剪区域（将图像从两边填充为正方形）。
|
|
|
|
|
*/
|
|
|
|
|
private fun initRectF(imageWidth: Int, imageHeight: Int): RectF {
|
|
|
|
|
val xMin: Float
|
|
|
|
|
@ -225,9 +227,8 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Checks whether there are enough torso keypoints.
|
|
|
|
|
* This function checks whether the model is confident at predicting one of the
|
|
|
|
|
* shoulders/hips which is required to determine a good crop region.
|
|
|
|
|
* 检查是否有足够的躯干关键点。
|
|
|
|
|
* 该函数检查模型是否能够自信地预测一个肩膀/臀部，这是确定良好裁剪区域所必需的。
|
|
|
|
|
*/
|
|
|
|
|
private fun torsoVisible(keyPoints: List<KeyPoint>): Boolean {
|
|
|
|
|
return ((keyPoints[BodyPart.LEFT_HIP.position].score > MIN_CROP_KEYPOINT_SCORE).or(
|
|
|
|
|
@ -240,13 +241,9 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Determines the region to crop the image for the model to run inference on.
|
|
|
|
|
* The algorithm uses the detected joints from the previous frame to estimate
|
|
|
|
|
* the square region that encloses the full body of the target person and
|
|
|
|
|
* centers at the midpoint of two hip joints. The crop size is determined by
|
|
|
|
|
* the distances between each joints and the center point.
|
|
|
|
|
* When the model is not confident with the four torso joint predictions, the
|
|
|
|
|
* function returns a default crop which is the full image padded to square.
|
|
|
|
|
* 根据上一帧的关键点来确定裁剪区域。
|
|
|
|
|
* 该算法使用上一帧检测到的关节来估算围绕目标人体的正方形区域，并以两个臀部关节的中点为中心。裁剪大小由每个关节与中心点的距离决定。
|
|
|
|
|
* 当模型对四个躯干关节的预测不自信时，函数返回默认的裁剪区域，即将整个图像填充为正方形。
|
|
|
|
|
*/
|
|
|
|
|
private fun determineRectF(
|
|
|
|
|
keyPoints: List<KeyPoint>,
|
|
|
|
|
@ -306,10 +303,9 @@ class MoveNet(private val interpreter: Interpreter, private var gpuDelegate: Gpu
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Calculates the maximum distance from each keypoints to the center location.
|
|
|
|
|
* The function returns the maximum distances from the two sets of keypoints:
|
|
|
|
|
* full 17 keypoints and 4 torso keypoints. The returned information will be
|
|
|
|
|
* used to determine the crop size. See determineRectF for more detail.
|
|
|
|
|
* * 计算每个关键点到中心位置的最大距离。
|
|
|
|
|
* * 该函数返回来自两组关键点的最大距离：
|
|
|
|
|
* * 全部17个关键点和4个躯干关键点。返回的信息将用于确定裁剪大小。
|
|
|
|
|
*/
|
|
|
|
|
private fun determineTorsoAndBodyDistances(
|
|
|
|
|
keyPoints: List<KeyPoint>,
|
|
|
|
|
|
