From 003b899023d6adf4f47b60a98c86fd2b9f7511f2 Mon Sep 17 00:00:00 2001
From: Xin Yang <xyang19@gmail.com>
Date: Mon, 10 Oct 2022 12:30:23 -0700
Subject: [PATCH] Add draw mask for BitMapWrapper

---
 android/build.gradle                          |  2 +-
 .../djl/android/core/BitmapImageFactory.java  | 56 ++++++++++++++++++-
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/android/build.gradle b/android/build.gradle
index cc13cf3c087..b2a50254244 100644
--- a/android/build.gradle
+++ b/android/build.gradle
@@ -6,7 +6,7 @@ buildscript {
     }
 
     dependencies {
-        classpath 'com.android.tools.build:gradle:4.2.2'
+        classpath 'com.android.tools.build:gradle:7.2.2'
     }
 }
 
diff --git a/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
index 87e2f64851f..7017370dd1b 100644
--- a/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
+++ b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
@@ -120,6 +120,13 @@ public Image fromNDArray(NDArray array) {
         return new BitMapWrapper(bitmap);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        Bitmap  bitmap = Bitmap.createBitmap(pixels, width, height, Bitmap.Config.ARGB_8888);
+        return new BitMapWrapper(bitmap);
+    }
+
     static class BitMapWrapper implements Image {
         private Bitmap bitmap;
 
@@ -141,10 +148,39 @@ public int getHeight() {
 
         /** {@inheritDoc} */
         @Override
-        public Object getWrappedImage() {
+        public Bitmap getWrappedImage() {
             return bitmap;
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public BitMapWrapper resize(int width, int height, boolean copy) {
+            if (!copy && bitmap.getWidth() == width && bitmap.getHeight() == height) {
+                return this;
+            }
+            return new BitMapWrapper(Bitmap.createScaledBitmap(bitmap, width, height, true));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Image getMask(int[][] mask) {
+            int w = mask[0].length;
+            int h = mask.length;
+            BitMapWrapper resized = resize(w, h, true);
+            Bitmap img = resized.getWrappedImage();
+            int[] pixels = new int[w * h];
+            int index = 0;
+            for (int y = 0; y < h; ++y) {
+                for (int x = 0; x < w; ++x) {
+                    if (mask[y][x] != 0) {
+                        pixels[index] = img.getPixel(x, y);
+                    }
+                    index++;
+                }
+            }
+            return new BitMapWrapper(Bitmap.createBitmap(pixels, w, h, Bitmap.Config.ARGB_8888));
+        }
+
         /** {@inheritDoc} */
         @Override
         public Image getSubImage(int x, int y, int w, int h) {
@@ -271,6 +307,24 @@ public void drawJoints(Joints joints) {
             oldBitmap.recycle();
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public void drawImage(Image overlay, boolean resize) {
+            if (!(overlay.getWrappedImage() instanceof Bitmap)) {
+                throw new IllegalArgumentException("Only Bitmap allowed");
+            }
+            if (resize) {
+                overlay = overlay.resize(getWidth(), getHeight(), false);
+            }
+            Bitmap target = Bitmap.createBitmap(getWidth(), getHeight(), Bitmap.Config.ARGB_8888);
+            Canvas canvas = new Canvas(target);
+            canvas.drawBitmap(bitmap, 0, 0, null);
+            canvas.drawBitmap((Bitmap) overlay.getWrappedImage(), 0, 0, null);
+            Bitmap oldBitmap = bitmap;
+            bitmap = target;
+            oldBitmap.recycle();
+        }
+
         private int randomColor() {
             return Color.rgb(
                     RandomUtils.nextInt(255), RandomUtils.nextInt(255), RandomUtils.nextInt(255));