Repository: deepjavalibrary/djl
Pull request #3595 (2025-01-23 21:54:02)

[fix] Fix the pytorch-native ci failure
## 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));
diff --git a/engines/pytorch/pytorch-native/build.gradle.kts b/engines/pytorch/pytorch-native/build.gradle.kts
index 917feed71..e5f002961 100644
--- a/engines/pytorch/pytorch-native/build.gradle.kts
+++ b/engines/pytorch/pytorch-native/build.gradle.kts
@@ -25,9 +25,16 @@ val isAarch64 = project.hasProperty("aarch64") || arch == "aarch64"
 version = ptVersion + if (isRelease) "" else "-SNAPSHOT"
 
 
-fun downloadBuild(ver: String, os: String, flavor: String, isPrecxx11: Boolean = false, isAarch64: Boolean = false) {
+fun downloadBuild(
+    execOperations: ExecOperations,
+    ver: String,
+    os: String,
+    flavor: String,
+    isPrecxx11: Boolean = false,
+    isAarch64: Boolean = false
+) {
     val arch = if (isAarch64) "aarch64" else "x86_64"
-    providers.exec {
+    execOperations.exec {
         workingDir = project.projectDir
         if (os == "win")
             commandLine(project.projectDir / "build.cmd", ver, flavor)
@@ -50,9 +57,9 @@ fun downloadBuild(ver: String, os: String, flavor: String, isPrecxx11: Boolean =
     }
 }
 
-fun downloadBuildAndroid(ver: String) {
+fun downloadBuildAndroid(execOperations: ExecOperations, ver: String) {
     for (abi in listOf("armeabi-v7a", "arm64-v8a", "x86", "x86_64")) {
-        providers.exec {
+        execOperations.exec {
             workingDir = project.projectDir
             commandLine("bash", "build_android.sh", ver, abi)
         }
@@ -65,7 +72,7 @@ fun downloadBuildAndroid(ver: String) {
     }
 }
 
-fun prepareNativeLib(binaryRoot: String, ver: String, packageType: String?) {
+fun prepareNativeLib(execOperations: ExecOperations, binaryRoot: String, ver: String, packageType: String?) {
     if ("mac" !in os)
         throw GradleException("This command must be run from osx")
 
@@ -96,7 +103,7 @@ fun prepareNativeLib(binaryRoot: String, ver: String, packageType: String?) {
         copyNativeLibToOutputDir(files, binaryRoot, officialPytorchUrl)
         copyNativeLibToOutputDir(aarch64Files, binaryRoot, aarch64PytorchUrl)
 
-        providers.exec {
+        execOperations.exec {
             workingDir = project.projectDir
             commandLine(
                 "install_name_tool",
@@ -105,7 +112,7 @@ fun prepareNativeLib(binaryRoot: String, ver: String, packageType: String?) {
                 "$binaryRoot/cpu/osx-aarch64/native/lib/libtorch_cpu.dylib"
             )
         }
-        providers.exec {
+        execOperations.exec {
             workingDir = project.projectDir
             commandLine(
                 "install_name_tool",
@@ -196,9 +203,9 @@ fun copyNativeLibToOutputDir(fileStoreMap: Map<String, String>, binaryRoot: Stri
 open class Cmd @Inject constructor(@Internal val execOperations: ExecOperations) : DefaultTask()
 
 tasks {
-    register("compileAndroidJNI") {
+    register<Cmd>("compileAndroidJNI") {
         doFirst {
-            downloadBuildAndroid(ptVersion)
+            downloadBuildAndroid(execOperations, ptVersion)
         }
     }
 
@@ -211,13 +218,13 @@ tasks {
         }
     }
 
-    register("compileJNI") {
+    register<Cmd>("compileJNI") {
         doFirst {
             // You have to use an environment with CUDA persets for Linux and Windows
             when {
-                "windows" in os -> downloadBuild(ptVersion, "win", ptFlavor)
-                "mac" in os -> downloadBuild(ptVersion, "osx", ptFlavor, false, isAarch64)
-                "linux" in os -> downloadBuild(ptVersion, "linux", ptFlavor, isPrecxx11, isAarch64)
+                "windows" in os -> downloadBuild(execOperations, ptVersion, "win", ptFlavor)
+                "mac" in os -> downloadBuild(execOperations, ptVersion, "osx", ptFlavor, false, isAarch64)
+                "linux" in os -> downloadBuild(execOperations, ptVersion, "linux", ptFlavor, isPrecxx11, isAarch64)
                 else -> throw IllegalStateException("Unknown Architecture $osName-$ptFlavor")
             }
 
@@ -228,10 +235,10 @@ tasks {
     }
 
     val binaryRoot = buildDirectory / "download"
-    register("downloadPyTorchNativeLib") {
+    register<Cmd>("downloadPyTorchNativeLib") {
         doLast {
             val packageType = project.findProperty("package_type")?.toString()
-            prepareNativeLib("$binaryRoot", ptVersion, packageType)
+            prepareNativeLib(execOperations, "$binaryRoot", ptVersion, packageType)
         }
     }
 
@@ -239,8 +246,8 @@ tasks {
         val dir = project.projectDir
         doLast {
             delete("$binaryRoot")
-            prepareNativeLib("$binaryRoot", ptVersion, "cpu")
-            prepareNativeLib("$binaryRoot", ptVersion, "gpu")
+            prepareNativeLib(execOperations, "$binaryRoot", ptVersion, "cpu")
+            prepareNativeLib(execOperations, "$binaryRoot", ptVersion, "gpu")
 
             execOperations.exec {
                 workingDir = dir


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);
+        }
+    }
+}