Internal Change

PiperOrigin-RevId: 399964034

Author: tensorflower-gardener

tensorflow_lite_support/odml/java/image/src/com/google/android/odml/image/BitmapExtractor.java

@@ -36,18 +36,14 @@ public final class BitmapExtractor {
    *     conversions.
    */
   public static Bitmap extract(MlImage image) {
-    ImageContainer container = image.getContainer();
-    switch (container.getImageProperties().getStorageType()) {
-      case MlImage.STORAGE_TYPE_BITMAP:
-        BitmapImageContainer bitmapImageContainer = (BitmapImageContainer) container;
-        return bitmapImageContainer.getBitmap();
-      case MlImage.STORAGE_TYPE_BYTEBUFFER:
-      case MlImage.STORAGE_TYPE_MEDIA_IMAGE:
-        // TODO(b/180504869): Support ByteBuffer -> Bitmap conversion.
-      default:
-        throw new IllegalArgumentException(
-            "Extracting Bitmap from an MlImage created by objects other than Bitmap is not"
-                + " supported");
+    ImageContainer imageContainer = image.getContainer(MlImage.STORAGE_TYPE_BITMAP);
+    if (imageContainer != null) {
+      return ((BitmapImageContainer) imageContainer).getBitmap();
+    } else {
+      // TODO(b/180504869): Support ByteBuffer -> Bitmap conversion.
+      throw new IllegalArgumentException(
+          "Extracting Bitmap from an MlImage created by objects other than Bitmap is not"
+              + " supported");
     }
   }
 

tensorflow_lite_support/odml/java/image/src/com/google/android/odml/image/ByteBufferExtractor.java

@@ -75,25 +75,31 @@ public static ByteBuffer extract(MlImage image) {
    *     conversions.
    */
   static ByteBuffer extract(MlImage image, @ImageFormat int targetFormat) {
-    ImageContainer container = image.getContainer();
-    switch (container.getImageProperties().getStorageType()) {
-      case MlImage.STORAGE_TYPE_BITMAP:
-        BitmapImageContainer bitmapImageContainer = (BitmapImageContainer) container;
-        return extractByteBufferFromBitmap(bitmapImageContainer.getBitmap(), targetFormat)
-            .asReadOnlyBuffer();
-      case MlImage.STORAGE_TYPE_BYTEBUFFER:
-        ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container;
-        @ImageFormat int sourceFormat = byteBufferImageContainer.getImageFormat();
-        if (sourceFormat == targetFormat) {
-          return byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer();
-        }
-        return convertByteBuffer(
-                byteBufferImageContainer.getByteBuffer(), sourceFormat, targetFormat)
-            .asReadOnlyBuffer();
-      default:
-        throw new IllegalArgumentException(
-            "Extracting ByteBuffer from an MlImage created by objects other than Bitmap or"
-                + " Bytebuffer is not supported");
+    ImageContainer container;
+    ImageProperties byteBufferProperties =
+        ImageProperties.builder()
+            .setStorageType(MlImage.STORAGE_TYPE_BYTEBUFFER)
+            .setImageFormat(targetFormat)
+            .build();
+    if ((container = image.getContainer(byteBufferProperties)) != null) {
+      ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container;
+      return byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer();
+    } else if ((container = image.getContainer(MlImage.STORAGE_TYPE_BYTEBUFFER)) != null) {
+      ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container;
+      @ImageFormat int sourceFormat = byteBufferImageContainer.getImageFormat();
+      return convertByteBuffer(byteBufferImageContainer.getByteBuffer(), sourceFormat, targetFormat)
+          .asReadOnlyBuffer();
+    } else if ((container = image.getContainer(MlImage.STORAGE_TYPE_BITMAP)) != null) {
+      BitmapImageContainer bitmapImageContainer = (BitmapImageContainer) container;
+      ByteBuffer byteBuffer =
+          extractByteBufferFromBitmap(bitmapImageContainer.getBitmap(), targetFormat)
+              .asReadOnlyBuffer();
+      image.addContainer(new ByteBufferImageContainer(byteBuffer, targetFormat));
+      return byteBuffer;
+    } else {
+      throw new IllegalArgumentException(
+          "Extracting ByteBuffer from an MlImage created by objects other than Bitmap or"
+              + " Bytebuffer is not supported");
     }
   }
 
@@ -128,23 +134,24 @@ static Result create(ByteBuffer buffer, @ImageFormat int imageFormat) {
    *     given {@code imageFormat}
    */
   static Result extractInRecommendedFormat(MlImage image) {
-    ImageContainer container = image.getContainer();
-    switch (container.getImageProperties().getStorageType()) {
-      case MlImage.STORAGE_TYPE_BITMAP:
-        BitmapImageContainer bitmapImageContainer = (BitmapImageContainer) container;
-        Bitmap bitmap = bitmapImageContainer.getBitmap();
-        @ImageFormat int format = adviseImageFormat(bitmap);
-        return Result.create(
-            extractByteBufferFromBitmap(bitmap, format).asReadOnlyBuffer(), format);
-      case MlImage.STORAGE_TYPE_BYTEBUFFER:
-        ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container;
-        return Result.create(
-            byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer(),
-            byteBufferImageContainer.getImageFormat());
-      default:
-        throw new IllegalArgumentException(
-            "Extract ByteBuffer from an MlImage created by objects other than Bitmap or Bytebuffer"
-                + " is not supported");
+    ImageContainer container;
+    if ((container = image.getContainer(MlImage.STORAGE_TYPE_BITMAP)) != null) {
+      Bitmap bitmap = ((BitmapImageContainer) container).getBitmap();
+      @ImageFormat int format = adviseImageFormat(bitmap);
+      Result result =
+          Result.create(extractByteBufferFromBitmap(bitmap, format).asReadOnlyBuffer(), format);
+
+      image.addContainer(new ByteBufferImageContainer(result.buffer(), result.format()));
+      return result;
+    } else if ((container = image.getContainer(MlImage.STORAGE_TYPE_BYTEBUFFER)) != null) {
+      ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container;
+      return Result.create(
+          byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer(),
+          byteBufferImageContainer.getImageFormat());
+    } else {
+      throw new IllegalArgumentException(
+          "Extract ByteBuffer from an MlImage created by objects other than Bitmap or Bytebuffer"
+              + " is not supported");
     }
   }
 

tensorflow_lite_support/odml/java/image/src/com/google/android/odml/image/ImageProperties.java

@@ -19,6 +19,7 @@
 import com.google.android.odml.image.MlImage.StorageType;
 import com.google.android.odml.image.annotation.KeepForSdk;
 import com.google.auto.value.AutoValue;
+import com.google.auto.value.extension.memoized.Memoized;
 
 /** Groups a set of properties to describe how an image is stored. */
 @AutoValue
@@ -40,6 +41,10 @@ public abstract class ImageProperties {
   @StorageType
   public abstract int getStorageType();
 
+  @Memoized
+  @Override
+  public abstract int hashCode();
+
   /**
    * Creates a builder of {@link ImageProperties}.
    *

tensorflow_lite_support/odml/java/image/src/com/google/android/odml/image/MediaImageExtractor.java

@@ -41,15 +41,12 @@ private MediaImageExtractor() {}
    * @throws IllegalArgumentException if the extraction failed.
    */
   public static Image extract(MlImage image) {
-    ImageContainer container = image.getContainer();
-    switch (container.getImageProperties().getStorageType()) {
-      case MlImage.STORAGE_TYPE_MEDIA_IMAGE:
-        MediaImageContainer mediaImageContainer = (MediaImageContainer) container;
-        return mediaImageContainer.getImage();
-      default:
-        throw new IllegalArgumentException(
-            "Extract Media Image from an MlImage created by objects other than Media Image"
-                + " is not supported");
+    ImageContainer container;
+    if ((container = image.getContainer(MlImage.STORAGE_TYPE_MEDIA_IMAGE)) != null) {
+      return ((MediaImageContainer) container).getImage();
     }
+    throw new IllegalArgumentException(
+        "Extract Media Image from an MlImage created by objects other than Media Image"
+            + " is not supported");
   }
 }

tensorflow_lite_support/odml/java/image/src/com/google/android/odml/image/MlImage.java

@@ -17,12 +17,16 @@
 
 import android.graphics.Rect;
 import androidx.annotation.IntDef;
+import androidx.annotation.Nullable;
 import com.google.android.odml.image.annotation.KeepForSdk;
 import java.io.Closeable;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 
 /**
  * Wraps image data for on-device machine learning (ODML) usages.
@@ -157,7 +161,9 @@ private synchronized void acquire() {
   public synchronized void close() {
     referenceCount -= 1;
     if (referenceCount == 0) {
-      getContainer().close();
+      for (ImageContainer imageContainer : containerMap.values()) {
+        imageContainer.close();
+      }
     }
   }
 
@@ -200,7 +206,7 @@ public Internal getInternal() {
     return new Internal(this);
   }
 
-  private final ImageContainer container;
+  private final Map<ImageProperties, ImageContainer> containerMap;
   private final int rotation;
   private final Rect roi;
   private final long timestamp;
@@ -212,7 +218,8 @@ public Internal getInternal() {
   /** Constructs an {@link MlImage} with a built container. */
   @KeepForSdk
   MlImage(ImageContainer container, int rotation, Rect roi, long timestamp, int width, int height) {
-    this.container = container;
+    this.containerMap = new HashMap<>();
+    containerMap.put(container.getImageProperties(), container);
     this.rotation = rotation;
     this.roi = new Rect();
     this.roi.set(roi);
@@ -232,7 +239,40 @@ ImageContainer getContainer() {
     // According to the design, in the future we will support multiple containers in one image.
     // Currently just return the original container.
     // TODO(b/182443927): Cache multiple containers in MlImage.
-    return container;
+    return containerMap.values().iterator().next();
+  }
+
+  /**
+   * Gets container from required {@code storageType}. Returns {@code null} if not existed.
+   *
+   * <p>If there are multiple containers with required {@code storageType}, returns the first one.
+   */
+  @Nullable
+  @KeepForSdk
+  ImageContainer getContainer(@StorageType int storageType) {
+    for (Entry<ImageProperties, ImageContainer> entry : containerMap.entrySet()) {
+      if (entry.getKey().getStorageType() == storageType) {
+        return entry.getValue();
+      }
+    }
+    return null;
+  }
+
+  /** Gets container from required {@code imageProperties}. Returns {@code null} if non existed. */
+  @Nullable
+  @KeepForSdk
+  ImageContainer getContainer(ImageProperties imageProperties) {
+    return containerMap.get(imageProperties);
+  }
+
+  /** Adds a new container if it doesn't exist. Returns {@code true} if it succeeds. */
+  boolean addContainer(ImageContainer container) {
+    ImageProperties imageProperties = container.getImageProperties();
+    if (containerMap.containsKey(imageProperties)) {
+      return false;
+    }
+    containerMap.put(imageProperties, container);
+    return true;
   }
 
   /**

tensorflow_lite_support/odml/java/image/tests/src/com/google/android/odml/image/BitmapMlImageBuilderTest.java

@@ -34,7 +34,7 @@ public void build_fromBitmap_succeeds() {
     Bitmap bitmap = Bitmap.createBitmap(20, 25, Config.ARGB_8888);
 
     MlImage image = new BitmapMlImageBuilder(bitmap).build();
-    ImageContainer container = image.getContainer();
+    ImageContainer container = image.getContainer(MlImage.STORAGE_TYPE_BITMAP);
 
     assertThat(image.getWidth()).isEqualTo(20);
     assertThat(image.getHeight()).isEqualTo(25);

tensorflow_lite_support/odml/java/image/tests/src/com/google/android/odml/image/ByteBufferExtractorTest.java

@@ -121,6 +121,16 @@ public void extract_rgbFromRgbaBitmap_succeeds() {
 
     assertThat(result.isReadOnly()).isTrue();
     assertThat(result).isEquivalentAccordingToCompareTo(TestImageCreator.createRgbBuffer());
+
+    // Verifies ByteBuffer is cached inside MlImage.
+    ByteBufferImageContainer byteBufferImageContainer =
+        (ByteBufferImageContainer) image.getContainer(MlImage.STORAGE_TYPE_BYTEBUFFER);
+    assertThat(byteBufferImageContainer.getByteBuffer()).isEqualTo(result);
+    assertThat(byteBufferImageContainer.getImageFormat()).isEqualTo(MlImage.IMAGE_FORMAT_RGB);
+
+    // Verifies that extracted ByteBuffer is the cached one.
+    ByteBuffer result2 = ByteBufferExtractor.extract(image, MlImage.IMAGE_FORMAT_RGB);
+    assertThat(result2).isEqualTo(result);
   }
 
   @Test
@@ -154,5 +164,16 @@ public void extractInRecommendedFormat_anyFormatFromRgbByteBuffer_succeeds() {
 
     assertThat(result.buffer().isReadOnly()).isTrue();
     assertThat(result.format()).isEqualTo(MlImage.IMAGE_FORMAT_RGB);
+
+    // Verifies ByteBuffer is cached inside MlImage.
+    ByteBufferImageContainer byteBufferImageContainer =
+        (ByteBufferImageContainer) image.getContainer(MlImage.STORAGE_TYPE_BYTEBUFFER);
+    assertThat(byteBufferImageContainer.getByteBuffer()).isEqualTo(result.buffer());
+    assertThat(byteBufferImageContainer.getImageFormat()).isEqualTo(MlImage.IMAGE_FORMAT_RGB);
+
+    // Verifies that extracted ByteBuffer is the cached one.
+    ByteBufferExtractor.Result result2 = ByteBufferExtractor.extractInRecommendedFormat(image);
+    assertThat(result2.buffer()).isEqualTo(result.buffer());
+    assertThat(result2.format()).isEqualTo(result.format());
   }
 }

tensorflow_lite_support/odml/java/image/tests/src/com/google/android/odml/image/ByteBufferMlImageBuilderTest.java

@@ -33,7 +33,7 @@ public void build_fromByteBuffer_succeeds() {
     ByteBuffer buffer = ByteBuffer.allocate(500);
 
     MlImage image = new ByteBufferMlImageBuilder(buffer, 20, 25, MlImage.IMAGE_FORMAT_RGB).build();
-    ImageContainer container = image.getContainer();
+    ImageContainer container = image.getContainer(MlImage.STORAGE_TYPE_BYTEBUFFER);
 
     assertThat(image.getWidth()).isEqualTo(20);
     assertThat(image.getHeight()).isEqualTo(25);

tensorflow_lite_support/odml/java/image/tests/src/com/google/android/odml/image/MediaMlImageBuilderTest.java

@@ -49,7 +49,7 @@ public void setUp() {
   @Test
   public void build_fromMediaImage_succeeds() {
     MlImage image = new MediaMlImageBuilder(mediaImage).build();
-    ImageContainer container = image.getContainer();
+    ImageContainer container = image.getContainer(MlImage.STORAGE_TYPE_MEDIA_IMAGE);
 
     assertThat(image.getWidth()).isEqualTo(WIDTH);
     assertThat(image.getHeight()).isEqualTo(HEIGHT);