[api] Adds serving support for some CV models (#3499)

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

Hide file tree

Showing 6 changed files with 24 additions and 67 deletions.
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
@@ -17,39 +17,26 @@
 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
@@ -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
@@ -16,33 +16,25 @@
 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;
  }
 }
diff --git a/...rc/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java b/...rc/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
@@ -14,7 +14,6 @@
 
 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 void setUp() {
 
  @Test
  public void testGetSupportedTypes() {
- Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+ Assert.assertEquals(factory.getSupportedTypes().size(), 6);
  }
 
  @Test
@@ -49,7 +48,7 @@ public void testNewInstance() {
  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
@@ -37,7 +37,7 @@ public void setUp() {
 
  @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
@@ -14,7 +14,6 @@
 
 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 void setUp() {
 
  @Test
  public void testGetSupportedTypes() {
- Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+ Assert.assertEquals(factory.getSupportedTypes().size(), 6);
  }
 
  @Test
@@ -48,7 +47,7 @@ public void testNewInstance() {
  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));
  }
  }
 }