Repository: deepjavalibrary/djl
Pull request #3593 (2025-01-23 17:35:19)

[api] Adds Pad image transform for CV
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/Pad.java b/api/src/main/java/ai/djl/modality/cv/transform/Pad.java
new file mode 100644
index 000000000..6768dbdd1
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/transform/Pad.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.cv.transform;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.Transform;
+
+/** A {@link Transform} that pad the image to square. */
+public class Pad implements Transform {
+
+    private double value;
+
+    /**
+     * Constructs a new {@code Pad} instance.
+     *
+     * @param value the padding value
+     */
+    public Pad(double value) {
+        this.value = value;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray transform(NDArray array) {
+        Shape shape = array.getShape();
+        int w = (int) shape.get(1);
+        int h = (int) shape.get(0);
+        if (w == h) {
+            return array;
+        }
+        int max = Math.max(w, h);
+        int padW = max - w;
+        int padH = max - h;
+        Shape padding = new Shape(0, 0, 0, padW, 0, padH);
+        array = array.pad(padding, value);
+        return array;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
index fb508e353..9f5beb81d 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
@@ -17,6 +17,7 @@ import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.transform.CenterCrop;
 import ai.djl.modality.cv.transform.CenterFit;
 import ai.djl.modality.cv.transform.Normalize;
+import ai.djl.modality.cv.transform.Pad;
 import ai.djl.modality.cv.transform.Resize;
 import ai.djl.modality.cv.transform.ResizeShort;
 import ai.djl.modality.cv.transform.ToTensor;
@@ -191,6 +192,14 @@ public abstract class BaseImageTranslator<T> implements Translator<Image, T> {
             if (arguments.containsKey("flag")) {
                 flag = Image.Flag.valueOf(arguments.get("flag").toString());
             }
+            String pad = ArgumentsUtil.stringValue(arguments, "pad", "false");
+            if ("true".equals(pad)) {
+                addTransform(new Pad(0));
+            } else if (!"false".equals(pad)) {
+                double padding = Double.parseDouble(pad);
+                addTransform(new Pad(padding));
+            }
+
             String resize = ArgumentsUtil.stringValue(arguments, "resize", "false");
             if ("true".equals(resize)) {
                 addTransform(new Resize(width, height));


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/cv/transform/PadTest.java b/api/src/test/java/ai/djl/modality/cv/transform/PadTest.java
new file mode 100644
index 000000000..c58dcf3fa
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/transform/PadTest.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.cv.transform;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class PadTest {
+
+    @Test
+    public void test() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(60, 50, 3));
+            Pad pad = new Pad(128);
+            NDArray output = pad.transform(array);
+            long[] shape = output.getShape().getShape();
+            Assert.assertEquals(shape[0], 60);
+            Assert.assertEquals(shape[1], 60);
+
+            array = manager.zeros(new Shape(50, 60, 3));
+            output = pad.transform(array);
+            shape = output.getShape().getShape();
+            Assert.assertEquals(shape[0], 60);
+            Assert.assertEquals(shape[1], 60);
+        }
+    }
+}