Repository: deepjavalibrary/djl
Pull request #1662 (2022-05-18 07:42:37)

[PaddleOCR] add openCV find rectangle method
## Description ##

Add OpenCV mechanism to do the rectangle findings



Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index 19956f893..e9af175ce 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -91,43 +91,28 @@ public class BufferedImageFactory extends ImageFactory {
         } else if (array.getDataType() != DataType.UINT8 && array.getDataType() != DataType.INT8) {
             throw new IllegalArgumentException("Datatype should be INT8 or UINT8");
         }
+        int height = (int) shape.get(0);
+        int width = (int) shape.get(1);
         if (NDImageUtils.isCHW(shape)) {
-            int height = (int) shape.get(1);
-            int width = (int) shape.get(2);
-            int imageArea = width * height;
-            BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
-            int[] raw = array.toUint8Array();
-            IntStream.range(0, imageArea)
-                    .forEach(
-                            ele -> {
-                                int x = ele % width;
-                                int y = ele / width;
-                                int red = raw[ele] & 0xFF;
-                                int green = raw[ele + imageArea] & 0xFF;
-                                int blue = raw[ele + imageArea * 2] & 0xFF;
-                                int rgb = (red << 16) | (green << 8) | blue;
-                                image.setRGB(x, y, rgb);
-                            });
-            return new BufferedImageWrapper(image);
-        } else {
-            int height = (int) shape.get(0);
-            int width = (int) shape.get(1);
-            int imageArea = width * height;
-            BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
-            int[] raw = array.toUint8Array();
-            IntStream.range(0, imageArea)
-                    .forEach(
-                            ele -> {
-                                int x = ele % width;
-                                int y = ele / width;
-                                int red = raw[ele * 3] & 0xFF;
-                                int green = raw[ele * 3 + 1] & 0xFF;
-                                int blue = raw[ele * 3 + 2] & 0xFF;
-                                int rgb = (red << 16) | (green << 8) | blue;
-                                image.setRGB(x, y, rgb);
-                            });
-            return new BufferedImageWrapper(image);
+            height = (int) shape.get(1);
+            width = (int) shape.get(2);
         }
+        int imageArea = width * height;
+        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
+        int[] raw = array.toUint8Array();
+        int finalWidth = width;
+        IntStream.range(0, imageArea)
+                .forEach(
+                        ele -> {
+                            int x = ele % finalWidth;
+                            int y = ele / finalWidth;
+                            int red = raw[ele] & 0xFF;
+                            int green = raw[ele + imageArea] & 0xFF;
+                            int blue = raw[ele + imageArea * 2] & 0xFF;
+                            int rgb = (red << 16) | (green << 8) | blue;
+                            image.setRGB(x, y, rgb);
+                        });
+        return new BufferedImageWrapper(image);
     }
 
     protected void save(BufferedImage image, OutputStream os, String type) throws IOException {
@@ -243,6 +228,13 @@ public class BufferedImageFactory extends ImageFactory {
             BufferedImageFactory.this.save(image, os, type);
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public List<BoundingBox> findBoundingBoxes() {
+            // TODO: Add grayscale conversion and use BoundFinder to implement
+            throw new UnsupportedOperationException("Not supported for BufferedImage");
+        }
+
         /** {@inheritDoc} */
         @Override
         public void drawBoundingBoxes(DetectedObjects detections) {
diff --git a/api/src/main/java/ai/djl/modality/cv/Image.java b/api/src/main/java/ai/djl/modality/cv/Image.java
index dea29361f..3053d85dd 100644
--- a/api/src/main/java/ai/djl/modality/cv/Image.java
+++ b/api/src/main/java/ai/djl/modality/cv/Image.java
@@ -12,12 +12,14 @@
  */
 package ai.djl.modality.cv;
 
+import ai.djl.modality.cv.output.BoundingBox;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.output.Joints;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.List;
 
 /**
  * {@code Image} is a container of an image in DJL. The storage type of the image depends on the
@@ -92,6 +94,13 @@ public interface Image {
      */
     void save(OutputStream os, String type) throws IOException;
 
+    /**
+     * Find bounding boxes from a masked image with 0/1 or 0/255.
+     *
+     * @return the List of bounding boxes of the images
+     */
+    List<BoundingBox> findBoundingBoxes();
+
     /**
      * Draws the bounding boxes on the image.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index c7c6f96aa..a48e56812 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -4683,7 +4683,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      *     href=https://d2l.djl.ai/chapter_linear-networks/softmax-regression.html#classification-problems>Classification-problems</a>
      */
     default NDArray oneHot(int depth, DataType dataType) {
-        return oneHot(depth, 0f, 1f, dataType);
+        return oneHot(depth, 1f, 0f, dataType);
     }
 
     /**
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle b/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
index d9ff67d4f..e18f1e7ed 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
@@ -3,7 +3,7 @@ group 'ai.djl.paddlepaddle'
 
 dependencies {
     api project(":engines:paddlepaddle:paddlepaddle-engine")
-
+    api project(":extensions:opencv")
     testImplementation("org.testng:testng:${testng_version}") {
         exclude group: "junit", module: "junit"
     }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslator.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslator.java
index 2f7ee2198..2312ef7e3 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslator.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslator.java
@@ -13,20 +13,25 @@
 package ai.djl.paddlepaddle.zoo.cv.objectdetection;
 
 import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.output.BoundingBox;
 import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.Rectangle;
 import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.opencv.OpenCVImageFactory;
 import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.NoBatchifyTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
+import java.io.FileNotFoundException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 /**
@@ -48,18 +53,37 @@ public class PpWordDetectionTranslator implements NoBatchifyTranslator<Image, De
 
     /** {@inheritDoc} */
     @Override
-    public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
+    public DetectedObjects processOutput(TranslatorContext ctx, NDList list)
+            throws FileNotFoundException {
         NDArray result = list.singletonOrThrow();
-        result = result.squeeze().mul(255f).toType(DataType.UINT8, true).neq(0);
-        boolean[] flattened = result.toBooleanArray();
-        Shape shape = result.getShape();
-        int w = (int) shape.get(0);
-        int h = (int) shape.get(1);
-        boolean[][] grid = new boolean[w][h];
-        IntStream.range(0, flattened.length)
-                .parallel()
-                .forEach(i -> grid[i / h][i % h] = flattened[i]);
-        List<BoundingBox> boxes = new BoundFinder(grid).getBoxes();
+        ImageFactory factory = ImageFactory.getInstance();
+        List<BoundingBox> boxes;
+        // faster mechanism
+        if (factory instanceof OpenCVImageFactory) {
+            result = result.squeeze(0).toType(DataType.UINT8, true);
+            Image image = factory.fromNDArray(result);
+            boxes =
+                    image.findBoundingBoxes()
+                            .parallelStream()
+                            .filter(
+                                    box -> {
+                                        Rectangle rect = (Rectangle) box;
+                                        return rect.getWidth() * image.getWidth() > 5
+                                                || rect.getHeight() * image.getHeight() > 5;
+                                    })
+                            .collect(Collectors.toList());
+        } else {
+            result = result.squeeze().mul(255f).toType(DataType.UINT8, true).neq(0);
+            boolean[] flattened = result.toBooleanArray();
+            Shape shape = result.getShape();
+            int w = (int) shape.get(0);
+            int h = (int) shape.get(1);
+            boolean[][] grid = new boolean[w][h];
+            IntStream.range(0, flattened.length)
+                    .parallel()
+                    .forEach(i -> grid[i / h][i % h] = flattened[i]);
+            boxes = new BoundFinder(grid).getBoxes();
+        }
         List<String> names = new ArrayList<>();
         List<Double> probs = new ArrayList<>();
         int boxSize = boxes.size();
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
index 5eb16744e..4453ebd3b 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
@@ -31,10 +31,13 @@ import java.awt.image.DataBufferByte;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.opencv.core.CvType;
 import org.opencv.core.Mat;
 import org.opencv.core.MatOfByte;
+import org.opencv.core.MatOfPoint;
 import org.opencv.core.Point;
 import org.opencv.core.Rect;
 import org.opencv.core.Scalar;
@@ -180,6 +183,25 @@ class OpenCVImage implements Image {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<BoundingBox> findBoundingBoxes() {
+        List<MatOfPoint> points = new ArrayList<>();
+        Imgproc.findContours(
+                image, points, new Mat(), Imgproc.RETR_LIST, Imgproc.CHAIN_APPROX_SIMPLE);
+        return points.parallelStream()
+                .map(
+                        point -> {
+                            Rect rect = Imgproc.boundingRect(point);
+                            return new Rectangle(
+                                    rect.x * 1.0 / image.width(),
+                                    rect.y * 1.0 / image.height(),
+                                    rect.width * 1.0 / image.width(),
+                                    rect.height * 1.0 / image.height());
+                        })
+                .collect(Collectors.toList());
+    }
+
     private void drawLandmarks(BoundingBox box) {
         Scalar color = new Scalar(0, 96, 246);
         for (ai.djl.modality.cv.output.Point point : box.getPath()) {
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
index f71f36b0a..7fbd4d4f6 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
@@ -74,11 +74,18 @@ public class OpenCVImageFactory extends ImageFactory {
         Shape shape = array.getShape();
         if (shape.dimension() == 4) {
             throw new UnsupportedOperationException("Batch is not supported");
-        } else if (shape.get(0) == 1 || shape.get(2) == 1) {
-            throw new UnsupportedOperationException("Grayscale image is not supported");
         } else if (array.getDataType() != DataType.UINT8 && array.getDataType() != DataType.INT8) {
             throw new IllegalArgumentException("Datatype should be INT8 or UINT8");
         }
+        boolean grayScale = shape.get(0) == 1 || shape.get(2) == 1;
+        if (grayScale) {
+            // expected CHW
+            int width = Math.toIntExact(shape.get(2));
+            int height = Math.toIntExact(shape.get(1));
+            Mat img = new Mat(height, width, CvType.CV_8UC1);
+            img.put(0, 0, array.toByteArray());
+            return new OpenCVImage(img);
+        }
         if (NDImageUtils.isCHW(shape)) {
             array = array.transpose(1, 2, 0);
             shape = array.getShape();


Test patch:
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/test/java/ai/djl/paddlepaddle/zoo/OCRTest.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/test/java/ai/djl/paddlepaddle/zoo/OCRTest.java
index 5bbffc49a..4941ea52a 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/test/java/ai/djl/paddlepaddle/zoo/OCRTest.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/test/java/ai/djl/paddlepaddle/zoo/OCRTest.java
@@ -17,6 +17,7 @@ import ai.djl.MalformedModelException;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Classifications;
+import ai.djl.modality.cv.BufferedImageFactory;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.output.BoundingBox;
@@ -39,7 +40,11 @@ public class OCRTest {
     public void testOCR() throws IOException, ModelException, TranslateException {
         String url = "https://resources.djl.ai/images/flight_ticket.jpg";
         Image img = ImageFactory.getInstance().fromUrl(url);
-        List<DetectedObjects.DetectedObject> boxes = detectWords(img).items();
+        DetectedObjects objects = detectWords(img);
+        if (ImageFactory.getInstance() instanceof BufferedImageFactory) {
+            objects = detectWords(img);
+        }
+        List<DetectedObjects.DetectedObject> boxes = objects.items();
 
         // load Character model
         Predictor<Image, String> recognizer = getRecognizer();
diff --git a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
index 009cb060f..9e7e47dfe 100644
--- a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
+++ b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
@@ -20,8 +20,10 @@ import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.output.Joints;
 import ai.djl.modality.cv.output.Landmark;
 import ai.djl.modality.cv.output.Mask;
+import ai.djl.modality.cv.output.Rectangle;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.testing.TestRequirements;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -29,6 +31,7 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import org.testng.Assert;
@@ -122,4 +125,31 @@ public class OpenCVImageFactoryTest {
                     });
         }
     }
+
+    @Test
+    public void testBoundingBoxes() {
+        TestRequirements.notWindows(); // failed on Windows ServerCore container
+        ImageFactory factory = ImageFactory.getInstance();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int[][] arr =
+                    new int[][] {
+                        {0, 1, 1, 1, 0},
+                        {0, 0, 1, 0, 0},
+                        {0, 0, 0, 0, 0},
+                        {1, 1, 0, 0, 0},
+                        {1, 0, 0, 0, 1}
+                    };
+            NDArray array = manager.create(arr).muli(255).toType(DataType.INT8, true).expandDims(0);
+            OpenCVImage image = (OpenCVImage) factory.fromNDArray(array);
+            List<BoundingBox> rectangles = image.findBoundingBoxes();
+            List<Rectangle> expected =
+                    Arrays.asList(
+                            new Rectangle(0.8, 0.8, 0.2, 0.2),
+                            new Rectangle(0, 0.6, 0.4, 0.4),
+                            new Rectangle(0.2, 0, 0.6, 0.4));
+            for (int i = 0; i < rectangles.size(); i++) {
+                Assert.assertEquals(rectangles.get(i).toString(), expected.get(i).toString());
+            }
+        }
+    }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
index a878ec1fe..df9b596a8 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
@@ -906,6 +906,7 @@ public class NDArrayOtherOpTest {
                     manager.create(
                             new float[][] {{0f, 1f, 0f}, {1f, 0f, 0f}, {0f, 0f, 1f}, {1f, 0f, 0f}});
             Assert.assertEquals(array.oneHot(3), expected);
+            Assert.assertEquals(array.oneHot(3, DataType.FLOAT32), expected);
             // test with all parameters
             array = manager.create(new int[] {1, 0, 2, 0});
             expected = manager.create(new int[][] {{1, 8, 1}, {8, 1, 1}, {1, 1, 8}, {8, 1, 1}});