From 714a6e555b106e7fd4de1b1e83d70e1c1c8570f3 Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Wed, 21 Dec 2022 08:06:20 -0800
Subject: [PATCH] Enable creating mediapipe image c++ packet directly from an
 Android media image object when its format is RGBA_8888.

PiperOrigin-RevId: 496923491
---
 .../mediapipe/framework/AndroidPacketCreator.java | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java b/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java
index 05700ba17..fc1e5484e 100644
--- a/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java
+++ b/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java
@@ -15,10 +15,13 @@
 package com.google.mediapipe.framework;
 
 import android.graphics.Bitmap;
+import android.graphics.PixelFormat;
+import android.media.Image;
 import com.google.mediapipe.framework.image.BitmapExtractor;
 import com.google.mediapipe.framework.image.ByteBufferExtractor;
 import com.google.mediapipe.framework.image.MPImage;
 import com.google.mediapipe.framework.image.MPImageProperties;
+import com.google.mediapipe.framework.image.MediaImageExtractor;
 import java.nio.ByteBuffer;
 
 // TODO: use Preconditions in this file.
@@ -97,7 +100,17 @@ public class AndroidPacketCreator extends PacketCreator {
       }
       return Packet.create(nativeCreateRgbaImage(mediapipeGraph.getNativeHandle(), bitmap));
     }
-
+    if (properties.getStorageType() == MPImage.STORAGE_TYPE_MEDIA_IMAGE) {
+      Image mediaImage = MediaImageExtractor.extract(image);
+      if (mediaImage.getFormat() != PixelFormat.RGBA_8888) {
+        throw new UnsupportedOperationException("Android media image must use RGBA_8888 config.");
+      }
+      return createImage(
+          mediaImage.getPlanes()[0].getBuffer(),
+          mediaImage.getWidth(),
+          mediaImage.getHeight(),
+          /* numChannels= */ 4);
+    }
     // Unsupported type.
     throw new UnsupportedOperationException(
         "Unsupported Image container type: " + properties.getStorageType());