Repository: deepjavalibrary/djl
Pull request #3499 (2024-10-06 17:47:30)

[api] Adds serving support for some CV models
## 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


codecov-commenter: :warning: Please install the !['codecov app svg image'](https://github.com/codecov/engineering-team/assets/152432831/e90313f4-9d3a-4b63-8b54-cfe14e7ec20d) to ensure uploads and comments are reliably processed by Codecov.

## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3499?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 58.92%. Comparing base [(`b6656ed`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/b6656ed60d823ee835bd69fb4e42a8244038b807?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`2c3487a`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/2c3487a6e85731f3d55da120db0639e99d6d864b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Report is 6 commits behind head on master.

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@              Coverage Diff              @@
##             master    #3499       +/-   ##
=============================================
- Coverage     72.52%   58.92%   -13.61%     
+ Complexity     7394     5851     -1543     
=============================================
  Files           672      674        +2     
  Lines         32909    33021      +112     
  Branches       3501     3516       +15     
=============================================
- Hits          23868    19457     -4411     
- Misses         7400    12082     +4682     
+ Partials       1641     1482      -159     
```

| [Flag](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3499/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3499/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | `?` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3499?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java b/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
index c0aea00b8..394a7461a 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
@@ -34,7 +34,9 @@ public class CategoryMask implements JsonSerializable {
 
     private static final int COLOR_BLACK = 0xFF000000;
 
+    @SuppressWarnings("serial")
     private List<String> classes;
+
     private int[][] mask;
 
     /**
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
index 751019c76..089d071c6 100644
--- 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.Translator;
 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 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
         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 class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
         }
 
         private static final class Prompt {
+
+            @SerializedName("image_url")
             String image;
+
             Location[] prompt;
             boolean visualize;
 
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
index 1bc7f5a5c..6840eca4d 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
@@ -17,39 +17,26 @@ import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.CategoryMask;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
 import java.io.Serializable;
-import java.lang.reflect.Type;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 
 /** A {@link TranslatorFactory} that creates a {@link SemanticSegmentationTranslator} instance. */
-public class SemanticSegmentationTranslatorFactory implements TranslatorFactory, Serializable {
+public class SemanticSegmentationTranslatorFactory extends BaseImageTranslatorFactory<CategoryMask>
+        implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
-    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
-
-    static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, CategoryMask.class));
-    }
-
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        if (input == Image.class && output == CategoryMask.class) {
-            return (Translator<I, O>) SemanticSegmentationTranslator.builder(arguments).build();
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, CategoryMask> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return SemanticSegmentationTranslator.builder(arguments).build();
     }
 
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return SUPPORTED_TYPES;
+    public Class<CategoryMask> getBaseOutputType() {
+        return CategoryMask.class;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
index c8fe84991..fd5696e5b 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
@@ -13,49 +13,29 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.Joints;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
 import java.io.Serializable;
-import java.lang.reflect.Type;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 
 /** An {@link TranslatorFactory} that creates a {@link SimplePoseTranslator} instance. */
-public class SimplePoseTranslatorFactory implements TranslatorFactory, Serializable {
+public class SimplePoseTranslatorFactory extends BaseImageTranslatorFactory<Joints>
+        implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
-    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
-
-    static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, Joints.class));
-        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
-    }
-
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return SUPPORTED_TYPES;
+    protected Translator<Image, Joints> buildBaseTranslator(Model model, Map<String, ?> arguments) {
+        return SimplePoseTranslator.builder(arguments).build();
     }
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        SimplePoseTranslator translator = SimplePoseTranslator.builder(arguments).build();
-        if (input == Image.class && output == Joints.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    public Class<Joints> getBaseOutputType() {
+        return Joints.class;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
index 4c492af50..a0c053210 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
@@ -16,33 +16,25 @@ import ai.djl.Model;
 import ai.djl.modality.cv.Image;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
 import java.io.Serializable;
-import java.lang.reflect.Type;
-import java.util.Collections;
 import java.util.Map;
-import java.util.Set;
 
 /** A {@link TranslatorFactory} that creates a {@link StyleTransferTranslator} instance. */
-public class StyleTransferTranslatorFactory implements TranslatorFactory, Serializable {
+public class StyleTransferTranslatorFactory extends BaseImageTranslatorFactory<Image>
+        implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return Collections.singleton(new Pair<>(Image.class, Image.class));
+    protected Translator<Image, Image> buildBaseTranslator(Model model, Map<String, ?> arguments) {
+        return new StyleTransferTranslator();
     }
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        if (!isSupported(input, output)) {
-            throw new IllegalArgumentException("Unsupported input/output types.");
-        }
-        return (Translator<I, O>) new StyleTransferTranslator();
+    public Class<Image> getBaseOutputType() {
+        return Image.class;
     }
 }


Test patch:
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
index a742e92d4..417896049 100644
--- 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 class Sam2InputTest {
         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"
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
index ec54e3bcb..9f2472950 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
@@ -14,7 +14,6 @@ package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
 import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.CategoryMask;
 import ai.djl.translate.Translator;
@@ -37,7 +36,7 @@ public class SemanticSegmentationTranslatorFactoryTest {
 
     @Test
     public void testGetSupportedTypes() {
-        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+        Assert.assertEquals(factory.getSupportedTypes().size(), 6);
     }
 
     @Test
@@ -49,7 +48,7 @@ public class SemanticSegmentationTranslatorFactoryTest {
             Assert.assertTrue(translator instanceof SemanticSegmentationTranslator);
             Assert.assertThrows(
                     IllegalArgumentException.class,
-                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+                    () -> factory.newInstance(Input.class, Image.class, model, arguments));
         }
     }
 }
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
index 0a2b54fd8..4eab9fa2e 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
@@ -37,7 +37,7 @@ public class SimplePoseTranslatorFactoryTest {
 
     @Test
     public void testGetSupportedTypes() {
-        Assert.assertEquals(factory.getSupportedTypes().size(), 2);
+        Assert.assertEquals(factory.getSupportedTypes().size(), 6);
     }
 
     @Test
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
index 785ffa872..1db9902b0 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
@@ -14,7 +14,6 @@ package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
 import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.translate.Translator;
 
@@ -36,7 +35,7 @@ public class StyleTransferTranslatorFactoryTest {
 
     @Test
     public void testGetSupportedTypes() {
-        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+        Assert.assertEquals(factory.getSupportedTypes().size(), 6);
     }
 
     @Test
@@ -48,7 +47,7 @@ public class StyleTransferTranslatorFactoryTest {
             Assert.assertTrue(translator instanceof StyleTransferTranslator);
             Assert.assertThrows(
                     IllegalArgumentException.class,
-                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+                    () -> factory.newInstance(Input.class, Image.class, model, arguments));
         }
     }
 }