Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 1 commit into from
Oct 6, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 7 additions & 20 deletions api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
32 changes: 6 additions & 26 deletions api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
20 changes: 6 additions & 14 deletions api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
5 changes: 2 additions & 3 deletions ...rc/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -37,7 +36,7 @@ public void setUp() {

@Test
public void testGetSupportedTypes() {
Assert.assertEquals(factory.getSupportedTypes().size(), 1);
Assert.assertEquals(factory.getSupportedTypes().size(), 6);
}

@Test
Expand All @@ -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));
}
}
}
2 changes: 1 addition & 1 deletion api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ public void setUp() {

@Test
public void testGetSupportedTypes() {
Assert.assertEquals(factory.getSupportedTypes().size(), 2);
Assert.assertEquals(factory.getSupportedTypes().size(), 6);
}

@Test
Expand Down
5 changes: 2 additions & 3 deletions api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -36,7 +35,7 @@ public void setUp() {

@Test
public void testGetSupportedTypes() {
Assert.assertEquals(factory.getSupportedTypes().size(), 1);
Assert.assertEquals(factory.getSupportedTypes().size(), 6);
}

@Test
Expand All @@ -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));
}
}
}
Loading