-
Notifications
You must be signed in to change notification settings - Fork 347
/
OverlayView.kt
118 lines (104 loc) · 4.02 KB
/
OverlayView.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
/*
* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.google.mediapipe.examples.poselandmarker
import android.content.Context
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.util.AttributeSet
import android.view.View
import androidx.core.content.ContextCompat
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 kotlin.math.max
import kotlin.math.min
class OverlayView(context: Context?, attrs: AttributeSet?) :
View(context, attrs) {
private var results: PoseLandmarkerResult? = null
private var pointPaint = Paint()
private var linePaint = Paint()
private var scaleFactor: Float = 1f
private var imageWidth: Int = 1
private var imageHeight: Int = 1
init {
initPaints()
}
fun clear() {
results = null
pointPaint.reset()
linePaint.reset()
invalidate()
initPaints()
}
private fun initPaints() {
linePaint.color =
ContextCompat.getColor(context!!, R.color.mp_color_primary)
linePaint.strokeWidth = LANDMARK_STROKE_WIDTH
linePaint.style = Paint.Style.STROKE
pointPaint.color = Color.YELLOW
pointPaint.strokeWidth = LANDMARK_STROKE_WIDTH
pointPaint.style = Paint.Style.FILL
}
override fun draw(canvas: Canvas) {
super.draw(canvas)
results?.let { poseLandmarkerResult ->
for(landmark in poseLandmarkerResult.landmarks()) {
for(normalizedLandmark in landmark) {
canvas.drawPoint(
normalizedLandmark.x() * imageWidth * scaleFactor,
normalizedLandmark.y() * imageHeight * scaleFactor,
pointPaint
)
}
PoseLandmarker.POSE_LANDMARKS.forEach {
canvas.drawLine(
poseLandmarkerResult.landmarks().get(0).get(it!!.start()).x() * imageWidth * scaleFactor,
poseLandmarkerResult.landmarks().get(0).get(it.start()).y() * imageHeight * scaleFactor,
poseLandmarkerResult.landmarks().get(0).get(it.end()).x() * imageWidth * scaleFactor,
poseLandmarkerResult.landmarks().get(0).get(it.end()).y() * imageHeight * scaleFactor,
linePaint)
}
}
}
}
fun setResults(
poseLandmarkerResults: PoseLandmarkerResult,
imageHeight: Int,
imageWidth: Int,
runningMode: RunningMode = RunningMode.IMAGE
) {
results = poseLandmarkerResults
this.imageHeight = imageHeight
this.imageWidth = imageWidth
scaleFactor = when (runningMode) {
RunningMode.IMAGE,
RunningMode.VIDEO -> {
min(width * 1f / imageWidth, height * 1f / imageHeight)
}
RunningMode.LIVE_STREAM -> {
// PreviewView is in FILL_START mode. So we need to scale up the
// landmarks to match with the size that the captured images will be
// displayed.
max(width * 1f / imageWidth, height * 1f / imageHeight)
}
}
invalidate()
}
companion object {
private const val LANDMARK_STROKE_WIDTH = 12F
}
}