Added code to draw lines.

PiperOrigin-RevId: 259806480

lite/examples/posenet/android/app/src/main/java/org/tensorflow/lite/examples/posenet/CameraConnectionFragment.kt

@@ -66,6 +66,8 @@ import java.util.concurrent.Semaphore
 import java.util.concurrent.TimeUnit
 import kotlin.collections.ArrayList
 import org.tensorflow.lite.Interpreter
+import org.tensorflow.lite.examples.posenet.lib.BodyPart
+import org.tensorflow.lite.examples.posenet.lib.Person
 import org.tensorflow.lite.examples.posenet.lib.Posenet
 
 class CameraConnectionFragment :
@@ -84,6 +86,22 @@ class CameraConnectionFragment :
     paint.setColor(Color.RED)
   }
 
+  /** List of body joints that should be connected.    */
+  val bodyJoints = listOf(
+    Pair(BodyPart.LEFT_WRIST, BodyPart.LEFT_ELBOW),
+    Pair(BodyPart.LEFT_ELBOW, BodyPart.LEFT_SHOULDER),
+    Pair(BodyPart.LEFT_SHOULDER, BodyPart.RIGHT_SHOULDER),
+    Pair(BodyPart.RIGHT_SHOULDER, BodyPart.RIGHT_ELBOW),
+    Pair(BodyPart.RIGHT_ELBOW, BodyPart.RIGHT_WRIST),
+    Pair(BodyPart.LEFT_SHOULDER, BodyPart.LEFT_HIP),
+    Pair(BodyPart.LEFT_HIP, BodyPart.RIGHT_HIP),
+    Pair(BodyPart.RIGHT_HIP, BodyPart.RIGHT_SHOULDER),
+    Pair(BodyPart.LEFT_HIP, BodyPart.LEFT_KNEE),
+    Pair(BodyPart.LEFT_KNEE, BodyPart.LEFT_ANKLE),
+    Pair(BodyPart.RIGHT_HIP, BodyPart.RIGHT_KNEE),
+    Pair(BodyPart.RIGHT_KNEE, BodyPart.RIGHT_ANKLE)
+  )
+
   /** An object for the Posenet library.    */
   private val posenet = Posenet()
 
@@ -507,14 +525,8 @@ class CameraConnectionFragment :
     }
   }
 
-  /** Process image using Posenet library.   */
-  private fun processImage(bitmap: Bitmap) {
-    // Create scaled version of bitmap for model input
-    var scaledBitmap = Bitmap.createScaledBitmap(bitmap, MODEL_WIDTH, MODEL_HEIGHT, true)
-    // Perform inference
-    val person = posenet.estimateSinglePose(interpreter!!, scaledBitmap)
-    var canvas: Canvas = surfaceHolder!!.lockCanvas()
-
+  /** Draw keypoints and lines.    */
+  private fun draw(canvas: Canvas, person: Person, bitmap: Bitmap) {
     // Draw bitmap on Canvas
     // TODO: crop bitmap to not squish it
     var screenWidth: Int = canvas.getWidth()
@@ -526,16 +538,44 @@ class CameraConnectionFragment :
       paint
     )
 
+    val widthRatio = screenWidth.toFloat() / MODEL_WIDTH
+    val heightRatio = screenHeight.toFloat() / MODEL_HEIGHT
+
     // Draw keypoints
     for (keypoint in person.keyPoints) {
-      var adjustedX: Float = keypoint.position.x.toFloat() / MODEL_WIDTH * screenWidth
-      var adjustedY: Float = keypoint.position.y.toFloat() / MODEL_HEIGHT * screenHeight
-      canvas.drawCircle(adjustedX, adjustedY, 8.0f, paint)
+      canvas.drawCircle(
+        keypoint.position.x.toFloat() * widthRatio,
+        keypoint.position.y.toFloat() * heightRatio,
+        8.0f, paint
+      )
+    }
+
+    paint.setStrokeWidth(8.0f)
+
+    for (line in person.bodyJoints) {
+      canvas.drawLine(
+        person.keyPoints[line.first.ordinal].position.x.toFloat() * widthRatio,
+        person.keyPoints[line.first.ordinal].position.y.toFloat() * heightRatio,
+        person.keyPoints[line.second.ordinal].position.x.toFloat() * widthRatio,
+        person.keyPoints[line.second.ordinal].position.y.toFloat() * heightRatio,
+        paint
+      )
     }
     // Draw!
     surfaceHolder!!.unlockCanvasAndPost(canvas)
   }
 
+  /** Process image using Posenet library.   */
+  private fun processImage(bitmap: Bitmap) {
+    // Create scaled version of bitmap for model input
+    var scaledBitmap = Bitmap.createScaledBitmap(bitmap, MODEL_WIDTH, MODEL_HEIGHT, true)
+
+    // Perform inference
+    val person = posenet.estimateSinglePose(interpreter!!, scaledBitmap)
+    var canvas: Canvas = surfaceHolder!!.lockCanvas()
+    draw(canvas, person, bitmap)
+  }
+
   /**
    * Creates a new [CameraCaptureSession] for camera preview.
    */

lite/examples/posenet/android/app/src/main/java/org/tensorflow/lite/examples/posenet/MainActivity.kt

@@ -32,10 +32,10 @@ import org.tensorflow.lite.examples.posenet.lib.Posenet as Posenet
 
 class MainActivity : AppCompatActivity() {
 
-  // Instantiate an Interpreter
+  /** Instantiate an Interpreter.    */
   private var interpreter: Interpreter? = null
 
-  // Preload and memory map the model file, returning a MappedByteBuffer containing the model.
+  /** Preload and memory map the model file, returns a MappedByteBuffer containing the model.    */
   fun loadModelFile(path: String): MappedByteBuffer {
     val fileDescriptor = assets.openFd(path)
     val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
@@ -44,7 +44,7 @@ class MainActivity : AppCompatActivity() {
     )
   }
 
-  // Returns a resized bitmap of the drawable image.
+  /** Returns a resized bitmap of the drawable image.    */
   private fun drawableToBitmap(drawable: Drawable): Bitmap {
     val bitmap = Bitmap.createBitmap(257, 353, Bitmap.Config.ARGB_8888)
     val canvas = Canvas(bitmap)
@@ -55,7 +55,7 @@ class MainActivity : AppCompatActivity() {
     return bitmap
   }
 
-  // Calls the Posenet library functions.
+  /** Calls the Posenet library functions.    */
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
     setContentView(R.layout.activity_main)
@@ -70,16 +70,16 @@ class MainActivity : AppCompatActivity() {
     val person = posenet.estimateSinglePose(interpreter!!, imageBitmap)
 
     // Draw the keypoints over the image.
-    val red = Paint()
-    red.setColor(Color.RED)
+    val paint = Paint()
+    paint.setColor(Color.RED)
     val size = 2.0f
 
     val mutableBitmap = imageBitmap.copy(Bitmap.Config.ARGB_8888, true)
     val canvas = Canvas(mutableBitmap)
-    for (i in 0 until person.keyPoints.size) {
+    for (keypoint in person.keyPoints) {
       canvas.drawCircle(
-        person.keyPoints[i].position.x.toFloat(),
-        person.keyPoints[i].position.y.toFloat(), size, red
+        keypoint.position.x.toFloat(),
+        keypoint.position.y.toFloat(), size, paint
       )
     }
     sampleImageView.setAdjustViewBounds(true)