[api] Makes Sam2 input consistent with other CV model (#3498)

deepjavalibrary · Oct 6, 2024 · 8be1c96 · 8be1c96
1 parent fea44e3
commit 8be1c96
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/Sam2Translator.java b/api/src/main/java/ai/djl/modality/cv/translator/Sam2Translator.java
@@ -39,6 +39,8 @@
 import ai.djl.translate.TranslatorContext;
 import ai.djl.util.JsonUtils;
 
+import com.google.gson.annotations.SerializedName;
+
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -340,7 +342,7 @@ float[][] getLabels() {
         public static Sam2Input fromJson(String input) throws IOException {
             Prompt prompt = JsonUtils.GSON.fromJson(input, Prompt.class);
             if (prompt.image == null) {
-                throw new IllegalArgumentException("Missing image value");
+                throw new IllegalArgumentException("Missing image_url value");
             }
             if (prompt.prompt == null || prompt.prompt.length == 0) {
                 throw new IllegalArgumentException("Missing prompt value");
@@ -477,7 +479,10 @@ public void setLabel(int label) {
         }
 
         private static final class Prompt {
+
+            @SerializedName("image_url")
             String image;
+
             Location[] prompt;
             boolean visualize;
 

diff --git a/api/src/test/java/ai/djl/modality/cv/translator/Sam2InputTest.java b/api/src/test/java/ai/djl/modality/cv/translator/Sam2InputTest.java
@@ -30,7 +30,7 @@ public void test() throws IOException {
         Path file = Paths.get("../examples/src/test/resources/kitten.jpg");
         Image img = ImageFactory.getInstance().fromFile(file);
         String json =
-                "{\"image\": \""
+                "{\"image_url\": \""
                         + file.toUri().toURL()
                         + "\",\n"
                         + "\"visualize\": true,\n"