From 3b13604ea48428841968ffcd0c100dcdf43d6a80 Mon Sep 17 00:00:00 2001
From: Xin Yang <xyang19@gmail.com>
Date: Sat, 8 Oct 2022 21:17:22 -0700
Subject: [PATCH] Add draw mask for BitMapWrapper

---
 android/build.gradle                          |  2 +-
 .../djl/android/core/BitmapImageFactory.java  | 53 ++++++++++++++++++-
 android/gradle.properties                     |  2 +-
 3 files changed, 54 insertions(+), 3 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..0b50ad80ccc 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,12 @@ public Image fromNDArray(NDArray array) {
         return new BitMapWrapper(bitmap);
     }
 
+    @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 +147,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 +306,22 @@ 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 = target;
+        }
+
         private int randomColor() {
             return Color.rgb(
                     RandomUtils.nextInt(255), RandomUtils.nextInt(255), RandomUtils.nextInt(255));
diff --git a/android/gradle.properties b/android/gradle.properties
index 5179882b54e..46e91c8cdd3 100644
--- a/android/gradle.properties
+++ b/android/gradle.properties
@@ -17,5 +17,5 @@ org.gradle.jvmargs=-Xmx1536m
 android.useAndroidX=true
 # Automatically convert third-party libraries to use AndroidX
 android.enableJetifier=true
-djl_version=0.19.0
+djl_version=0.20.0
 pytorch_version=1.12.1