Repository: deepjavalibrary/djl
Pull request #2032 (2022-09-21 17:39:03)

update the SemanticSegmentationTranslator

## Description ##

Brief description of what this PR is about

update the SemanticSegmentationTranslator to get the original photo pixels


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2032?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **69.64**% // Decreases project coverage by **`-2.44%`** :warning:
> Coverage data is based on head [(`1fd55ad`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2032?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 68.38% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2032      +/-   ##
============================================
- Coverage     72.08%   69.64%   -2.45%     
- Complexity     5126     5954     +828     
============================================
  Files           473      596     +123     
  Lines         21970    26463    +4493     
  Branches       2351     2876     +525     
============================================
+ Hits          15838    18429    +2591     
- Misses         4925     6651    +1726     
- Partials       1207     1383     +176     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2032?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [.../cv/translator/InstanceSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (-86.59%)` | :arrow_down: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `7.14% <0.00%> (-11.04%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| ... and [510 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2032/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2032?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](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/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index a66050657..d850b08a1 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -127,6 +127,14 @@ public class BufferedImageFactory extends ImageFactory {
         return new BufferedImageWrapper(image);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
+        image.setRGB(0, 0, width, height, pixels, 0, width);
+        return new BufferedImageWrapper(image);
+    }
+
     protected void save(BufferedImage image, OutputStream os, String type) throws IOException {
         ImageIO.write(image, type, os);
     }
@@ -153,10 +161,26 @@ public class BufferedImageFactory extends ImageFactory {
 
         /** {@inheritDoc} */
         @Override
-        public Object getWrappedImage() {
+        public BufferedImage getWrappedImage() {
             return image;
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public BufferedImageWrapper resize(int width, int height, boolean copy) {
+            if (!copy && image.getWidth() == width && image.getHeight() == height) {
+                return this;
+            }
+
+            java.awt.Image img =
+                    image.getScaledInstance(width, height, java.awt.Image.SCALE_SMOOTH);
+            BufferedImage scaled = new BufferedImage(width, height, image.getType());
+            Graphics2D g2d = scaled.createGraphics();
+            g2d.drawImage(img, 0, 0, null);
+            g2d.dispose();
+            return new BufferedImageWrapper(scaled);
+        }
+
         /** {@inheritDoc} */
         @Override
         public Image getSubImage(int x, int y, int w, int h) {
@@ -174,6 +198,26 @@ public class BufferedImageFactory extends ImageFactory {
             return new BufferedImageWrapper(copy);
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public Image getMask(int[][] mask) {
+            int w = mask[0].length;
+            int h = mask.length;
+            BufferedImageWrapper resized = resize(w, h, true);
+            BufferedImage img = resized.getWrappedImage();
+            int[] pixels = new int[w * h];
+            int index = 0;
+            for (int y = 0; y < h; ++y) {
+                for (int x = 0; x < w; ++x) {
+                    if (mask[y][x] != 0) {
+                        pixels[index] = img.getRGB(x, y);
+                    }
+                    index++;
+                }
+            }
+            return fromPixels(pixels, w, h);
+        }
+
         private void convertIdNeeded() {
             if (image.getType() == BufferedImage.TYPE_INT_ARGB) {
                 return;
@@ -308,6 +352,24 @@ public class BufferedImageFactory extends ImageFactory {
             g.dispose();
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public void drawImage(Image overlay, boolean resize) {
+            if (!(overlay.getWrappedImage() instanceof BufferedImage)) {
+                throw new IllegalArgumentException("Only BufferedImage allowed");
+            }
+            if (resize) {
+                overlay = overlay.resize(getWidth(), getHeight(), false);
+            }
+            BufferedImage target =
+                    new BufferedImage(getWidth(), getHeight(), BufferedImage.TYPE_INT_ARGB);
+            Graphics2D g = (Graphics2D) target.getGraphics();
+            g.drawImage(image, 0, 0, null);
+            g.drawImage((BufferedImage) overlay.getWrappedImage(), 0, 0, null);
+            g.dispose();
+            image = target;
+        }
+
         private Color randomColor() {
             return new Color(RandomUtils.nextInt(255));
         }
diff --git a/api/src/main/java/ai/djl/modality/cv/Image.java b/api/src/main/java/ai/djl/modality/cv/Image.java
index 6e93aba69..26e7fe4ce 100644
--- a/api/src/main/java/ai/djl/modality/cv/Image.java
+++ b/api/src/main/java/ai/djl/modality/cv/Image.java
@@ -49,6 +49,24 @@ public interface Image {
      */
     Object getWrappedImage();
 
+    /**
+     * Creates a new resized image.
+     *
+     * @param width the new image's desired width
+     * @param height the new image's desired height
+     * @param copy false to return original image if size is the same
+     * @return the new resized image.
+     */
+    Image resize(int width, int height, boolean copy);
+
+    /**
+     * Returns a new {@code Image} of masked area.
+     *
+     * @param mask the mask for each pixel in the image
+     * @return the mask image.
+     */
+    Image getMask(int[][] mask);
+
     /**
      * Gets the subimage defined by a specified rectangular region.
      *
@@ -116,6 +134,14 @@ public interface Image {
      */
     void drawJoints(Joints joints);
 
+    /**
+     * Draws the overlay on the image.
+     *
+     * @param overlay the overlay image
+     * @param resize true to resize the overlay image to match the image
+     */
+    void drawImage(Image overlay, boolean resize);
+
     /** Flag indicates the color channel options for images. */
     enum Flag {
         GRAYSCALE,
diff --git a/api/src/main/java/ai/djl/modality/cv/ImageFactory.java b/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
index 5e73d7ae6..17167a8b6 100644
--- a/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
@@ -138,4 +138,14 @@ public abstract class ImageFactory {
      * @return {@link Image}
      */
     public abstract Image fromNDArray(NDArray array);
+
+    /**
+     * Gets {@link Image} from array.
+     *
+     * @param pixels the array of ARGB values used to initialize the pixels.
+     * @param width the width of the image
+     * @param height the height of the image
+     * @return {@link Image}
+     */
+    public abstract Image fromPixels(int[] pixels, int width, int height);
 }
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
new file mode 100644
index 000000000..15ec37c75
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
@@ -0,0 +1,205 @@
+/*
+ * Copyright 2022 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.output;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.util.JsonSerializable;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.RandomUtils;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonSerializationContext;
+import com.google.gson.JsonSerializer;
+
+import java.awt.Color;
+import java.lang.reflect.Type;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+/**
+ * A class representing the segmentation result of an image in an {@link
+ * ai.djl.Application.CV#SEMANTIC_SEGMENTATION} case.
+ */
+public class CategoryMask implements JsonSerializable {
+
+    private static final long serialVersionUID = 1L;
+
+    private static final Gson GSON =
+            JsonUtils.builder()
+                    .registerTypeAdapter(CategoryMask.class, new SegmentationSerializer())
+                    .create();
+
+    private transient List<String> classes;
+    private int[][] mask;
+
+    /**
+     * Constructs a Mask with the given data.
+     *
+     * @param classes the list of classes
+     * @param mask the category mask for each pixel in the image
+     */
+    public CategoryMask(List<String> classes, int[][] mask) {
+        this.classes = classes;
+        this.mask = mask;
+    }
+
+    /**
+     * Returns the list of classes.
+     *
+     * @return list of classes
+     */
+    public List<String> getClasses() {
+        return classes;
+    }
+
+    /**
+     * Returns the class for each pixel.
+     *
+     * @return the class for each pixel
+     */
+    public int[][] getMask() {
+        return mask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ByteBuffer toByteBuffer() {
+        return ByteBuffer.wrap(toJson().getBytes(StandardCharsets.UTF_8));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String toJson() {
+        return GSON.toJson(this) + '\n';
+    }
+
+    /**
+     * Extracts the detected objects from the image.
+     *
+     * @param image the original image
+     * @return the detected objects from the image
+     */
+    public Image getMaskImage(Image image) {
+        return image.getMask(mask);
+    }
+
+    /**
+     * Extracts the specified object from the image.
+     *
+     * @param image the original image
+     * @param classId the class to extract from the image
+     * @return the specific object on the image
+     */
+    public Image getMaskImage(Image image, int classId) {
+        int width = mask[0].length;
+        int height = mask.length;
+        int[][] selected = new int[height][width];
+        for (int h = 0; h < height; ++h) {
+            for (int w = 0; w < width; ++w) {
+                selected[h][w] = mask[h][w] == classId ? 1 : 0;
+            }
+        }
+        return image.getMask(selected);
+    }
+
+    /**
+     * Extracts the background from the image.
+     *
+     * @param image the original image
+     * @return the background of the image
+     */
+    public Image getBackgroundImage(Image image) {
+        return getMaskImage(image, 0);
+    }
+
+    /**
+     * Highlights the detected object on the image with random colors.
+     *
+     * @param image the original image
+     * @param transparency the transparency of the overlay. Value is between 0 and 255 inclusive,
+     *     where 0 means the overlay is completely transparent and 255 means the overlay is
+     *     completely opaque.
+     */
+    public void drawMask(Image image, int transparency) {
+        drawMask(image, transparency, Color.BLACK.getRGB());
+    }
+
+    /**
+     * Highlights the detected object on the image with random colors.
+     *
+     * @param image the original image
+     * @param transparency the transparency of the overlay. Value is between 0 and 255 inclusive,
+     *     where 0 means the overlay is completely opaque and 255 means the overlay is completely
+     *     transparent.
+     * @param background replace the background with specified background color, use transparent
+     *     color to remove background
+     */
+    public void drawMask(Image image, int transparency, int background) {
+        int[] colors = generateColors(background, transparency);
+        Image maskImage = getColorOverlay(colors);
+        image.drawImage(maskImage, true);
+    }
+
+    /**
+     * Highlights the specified object with specific color.
+     *
+     * @param image the original image
+     * @param classId the class to draw on the image
+     * @param color the rgb color with transparency
+     * @param transparency the transparency of the overlay
+     */
+    public void drawMask(Image image, int classId, int color, int transparency) {
+        int[] colors = new int[classes.size()];
+        colors[classId] = color & 0xFFFFFF | transparency << 24;
+        Image colorOverlay = getColorOverlay(colors);
+        image.drawImage(colorOverlay, true);
+    }
+
+    private Image getColorOverlay(int[] colors) {
+        int height = mask.length;
+        int width = mask[0].length;
+        int[] pixels = new int[width * height];
+        for (int h = 0; h < height; h++) {
+            for (int w = 0; w < width; w++) {
+                int index = mask[h][w];
+                pixels[h * width + w] = colors[index];
+            }
+        }
+        return ImageFactory.getInstance().fromPixels(pixels, width, height);
+    }
+
+    private int[] generateColors(int background, int transparency) {
+        int[] colors = new int[classes.size()];
+        colors[0] = background;
+        for (int i = 1; i < classes.size(); i++) {
+            int red = RandomUtils.nextInt(256);
+            int green = RandomUtils.nextInt(256);
+            int blue = RandomUtils.nextInt(256);
+            colors[i] = transparency << 24 | red << 16 | green << 8 | blue;
+        }
+        return colors;
+    }
+
+    /** A customized Gson serializer to serialize the {@code Segmentation} object. */
+    public static final class SegmentationSerializer implements JsonSerializer<CategoryMask> {
+
+        /** {@inheritDoc} */
+        @Override
+        public JsonElement serialize(CategoryMask src, Type type, JsonSerializationContext ctx) {
+            return ctx.serialize(src.getMask());
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
index b9e0b5726..9ca5942a7 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
@@ -13,33 +13,31 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.ImageFactory;
-import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.CategoryMask;
 import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.Transform;
+import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
-import ai.djl.util.RandomUtils;
 
-import java.nio.ByteBuffer;
+import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 /**
- * A {@link BaseImageTranslator} that post-process the {@link NDArray} into {@link DetectedObjects}
- * with boundaries at the detailed pixel level.
+ * A {@link Translator} that post-process the {@link Image} into {@link CategoryMask} with output
+ * mask representing the class that each pixel in the original image belong to.
  */
-public class SemanticSegmentationTranslator extends BaseImageTranslator<Image> {
+public class SemanticSegmentationTranslator extends BaseImageTranslator<CategoryMask> {
 
+    private SynsetLoader synsetLoader;
     private final int shortEdge;
     private final int maxEdge;
 
-    private static final int CHANNEL = 3;
-    private static final int CLASSNUM = 21;
+    private List<String> classes;
 
     /**
      * Creates the Semantic Segmentation translator from the given builder.
@@ -48,78 +46,58 @@ public class SemanticSegmentationTranslator extends BaseImageTranslator<Image> {
      */
     public SemanticSegmentationTranslator(Builder builder) {
         super(builder);
+        this.synsetLoader = builder.synsetLoader;
         this.shortEdge = builder.shortEdge;
         this.maxEdge = builder.maxEdge;
 
         pipeline.insert(0, null, new ResizeShort());
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        if (classes == null) {
+            classes = synsetLoader.load(ctx.getModel());
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDList processInput(TranslatorContext ctx, Image image) {
-        ctx.setAttachment("originalHeight", image.getHeight());
-        ctx.setAttachment("originalWidth", image.getWidth());
         return super.processInput(ctx, image);
     }
 
     /** {@inheritDoc} */
     @Override
-    public Image processOutput(TranslatorContext ctx, NDList list) {
+    public CategoryMask processOutput(TranslatorContext ctx, NDList list) {
         // scores contains the probabilities of each pixel being a certain object
         float[] scores = list.get(1).toFloatArray();
         Shape shape = list.get(1).getShape();
         int width = (int) shape.get(2);
         int height = (int) shape.get(1);
-
-        // build image array
-        try (NDManager manager = NDManager.newBaseManager()) {
-            int imageSize = width * height;
-            ByteBuffer bb = manager.allocateDirect(CHANNEL * imageSize);
-            int r = 0; // adjustment for red pixel
-            int g = 1; // adjustment for green pixel
-            int b = 2; // adjustment for blue pixel
-            byte[][] colors = new byte[CLASSNUM][3];
-            for (int i = 0; i < CLASSNUM; i++) {
-                byte red = (byte) RandomUtils.nextInt(256);
-                byte green = (byte) RandomUtils.nextInt(256);
-                byte blue = (byte) RandomUtils.nextInt(256);
-                colors[i][r] = red;
-                colors[i][g] = green;
-                colors[i][b] = blue;
-            }
-
-            // change color of pixels in image array where objects have been detected
-            for (int h = 0; h < height; h++) {
-                for (int w = 0; w < width; w++) {
-                    int index = h * width + w;
-                    int maxi = 0;
-                    double maxnum = -Double.MAX_VALUE;
-                    for (int i = 0; i < CLASSNUM; i++) {
-                        // get score for each i at the h,w pixel of the image
-                        float score = scores[i * (imageSize) + index];
-                        if (score > maxnum) {
-                            maxnum = score;
-                            maxi = i;
-                        }
-                    }
-                    if (maxi > 0) {
-                        bb.put(colors[maxi][r]);
-                        bb.put(colors[maxi][g]);
-                        bb.put(colors[maxi][b]);
-                    } else {
-                        bb.position(bb.position() + 3);
+        int[][] mask = new int[height][width];
+
+        int imageSize = width * height;
+
+        // Build mask array
+        int numOfClasses = classes.size();
+        for (int h = 0; h < height; h++) {
+            for (int w = 0; w < width; w++) {
+                int index = h * width + w;
+                int maxi = 0;
+                double maxnum = -Double.MAX_VALUE;
+                for (int i = 0; i < numOfClasses; ++i) {
+                    // get score for each i at the h,w pixel of the image
+                    float score = scores[i * imageSize + index];
+                    if (score > maxnum) {
+                        maxnum = score;
+                        maxi = i;
                     }
                 }
+                mask[h][w] = maxi;
             }
-            bb.rewind();
-            int originW = (int) ctx.getAttachment("originalWidth");
-            int originH = (int) ctx.getAttachment("originalHeight");
-            NDArray fullImage =
-                    manager.create(bb, new Shape(height, width, CHANNEL), DataType.UINT8);
-            NDArray resized = NDImageUtils.resize(fullImage, originW, originH);
-
-            return ImageFactory.getInstance().fromNDArray(resized);
         }
+        return new CategoryMask(classes, mask);
     }
 
     /**
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 ec69dabdd..6374fc1fa 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
@@ -13,20 +13,13 @@
 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.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+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.InputStream;
 import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -37,7 +30,7 @@ public class SemanticSegmentationTranslatorFactory implements TranslatorFactory
     private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
 
     static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, Image.class));
+        SUPPORTED_TYPES.add(new Pair<>(Image.class, CategoryMask.class));
     }
 
     /** {@inheritDoc} */
@@ -45,18 +38,8 @@ public class SemanticSegmentationTranslatorFactory implements TranslatorFactory
     @SuppressWarnings("unchecked")
     public <I, O> Translator<I, O> newInstance(
             Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        SemanticSegmentationTranslator translator =
-                SemanticSegmentationTranslator.builder(arguments).build();
-        if (input == Image.class && output == Image.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == Image.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == Image.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == Image.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
+        if (input == Image.class && output == CategoryMask.class) {
+            return (Translator<I, O>) SemanticSegmentationTranslator.builder(arguments).build();
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
index 92a677484..ef281e2c0 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
@@ -12,14 +12,12 @@
  */
 package ai.djl.examples.inference;
 
-import ai.djl.Application;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
-import ai.djl.modality.cv.transform.Normalize;
-import ai.djl.modality.cv.transform.ToTensor;
-import ai.djl.modality.cv.translator.SemanticSegmentationTranslator;
+import ai.djl.modality.cv.output.CategoryMask;
+import ai.djl.modality.cv.translator.SemanticSegmentationTranslatorFactory;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.util.ProgressBar;
@@ -28,6 +26,7 @@ import ai.djl.translate.TranslateException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.awt.Color;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -42,8 +41,7 @@ import java.nio.file.Paths;
  */
 public final class SemanticSegmentation {
 
-    private static final Logger logger =
-            LoggerFactory.getLogger(SemanticSegmentationTranslator.class);
+    private static final Logger logger = LoggerFactory.getLogger(SemanticSegmentation.class);
 
     private SemanticSegmentation() {}
 
@@ -53,41 +51,58 @@ public final class SemanticSegmentation {
 
     public static void predict() throws IOException, ModelException, TranslateException {
         Path imageFile = Paths.get("src/test/resources/dog_bike_car.jpg");
-        Image img = ImageFactory.getInstance().fromFile(imageFile);
-
-        float[] mean = {0.485f, 0.456f, 0.406f};
-        float[] std = {0.229f, 0.224f, 0.225f};
+        ImageFactory factory = ImageFactory.getInstance();
+        Image img = factory.fromFile(imageFile);
 
         String url =
                 "https://mlrepo.djl.ai/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/0.0.1/deeplabv3.zip";
-        SemanticSegmentationTranslator translator =
-                SemanticSegmentationTranslator.builder()
-                        .addTransform(new ToTensor())
-                        .addTransform(new Normalize(mean, std))
-                        .build();
 
-        Criteria<Image, Image> criteria =
+        Criteria<Image, CategoryMask> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.SEMANTIC_SEGMENTATION)
-                        .setTypes(Image.class, Image.class)
+                        .setTypes(Image.class, CategoryMask.class)
                         .optModelUrls(url)
-                        .optTranslator(translator)
+                        .optTranslatorFactory(new SemanticSegmentationTranslatorFactory())
                         .optEngine("PyTorch")
                         .optProgress(new ProgressBar())
                         .build();
-        try (ZooModel<Image, Image> model = criteria.loadModel();
-                Predictor<Image, Image> predictor = model.newPredictor()) {
-            Image semanticImage = predictor.predict(img);
-            saveSemanticImage(semanticImage);
+        Image bg = factory.fromFile(Paths.get("../website/img/stars-in-the-night-sky.jpg"));
+        try (ZooModel<Image, CategoryMask> model = criteria.loadModel();
+                Predictor<Image, CategoryMask> predictor = model.newPredictor()) {
+            CategoryMask mask = predictor.predict(img);
+
+            // Highlights the detected object on the image with random opaque colors.
+            Image img1 = img.duplicate();
+            mask.drawMask(img1, 255);
+            saveSemanticImage(img1, "semantic_instances1.png");
+
+            // Highlights the detected object on the image with random colors.
+            Image img2 = img.duplicate();
+            mask.drawMask(img2, 180, 0);
+            saveSemanticImage(img2, "semantic_instances2.png");
+
+            // Highlights only the dog with blue color.
+            Image img3 = img.duplicate();
+            mask.drawMask(img3, 12, Color.BLUE.brighter().getRGB(), 180);
+            saveSemanticImage(img3, "semantic_instances3.png");
+
+            // Extract dog from the image
+            Image dog = mask.getMaskImage(img, 12);
+            dog = dog.resize(img.getWidth(), img.getHeight(), true);
+            saveSemanticImage(dog, "semantic_instances4.png");
+
+            // Replace background with an image
+            bg = bg.resize(img.getWidth(), img.getHeight(), true);
+            bg.drawImage(dog, true);
+            saveSemanticImage(bg, "semantic_instances5.png");
         }
     }
 
-    private static void saveSemanticImage(Image img) throws IOException {
+    private static void saveSemanticImage(Image img, String fileName) throws IOException {
         Path outputDir = Paths.get("build/output");
         Files.createDirectories(outputDir);
 
-        Path imagePath = outputDir.resolve("semantic_instances.png");
+        Path imagePath = outputDir.resolve(fileName);
         img.save(Files.newOutputStream(imagePath), "png");
-        logger.info("Segmentation result image has been saved in: {}", imagePath);
+        logger.info("Segmentation result image has been saved in: {}", imagePath.toAbsolutePath());
     }
 }
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
index 9aac8e799..ebcb9b843 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
@@ -39,7 +39,9 @@ import org.opencv.imgproc.Imgproc;
 import java.awt.Color;
 import java.awt.Graphics2D;
 import java.awt.image.BufferedImage;
+import java.awt.image.DataBuffer;
 import java.awt.image.DataBufferByte;
+import java.awt.image.DataBufferInt;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
@@ -75,10 +77,41 @@ class OpenCVImage implements Image {
 
     /** {@inheritDoc} */
     @Override
-    public Object getWrappedImage() {
+    public Mat getWrappedImage() {
         return image;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public OpenCVImage resize(int width, int height, boolean copy) {
+        if (!copy && image.width() == width && image.height() == height) {
+            return this;
+        }
+
+        Mat resized = new Mat();
+        Imgproc.resize(image, resized, new Size(width, height));
+        return new OpenCVImage(resized);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Image getMask(int[][] mask) {
+        int w = mask[0].length;
+        int h = mask.length;
+        OpenCVImage resized = resize(w, h, false);
+        Mat img = resized.getWrappedImage();
+        Mat ret = new Mat(h, w, CvType.CV_8UC4);
+        for (int y = 0; y < h; ++y) {
+            for (int x = 0; x < w; ++x) {
+                if (mask[y][x] != 0) {
+                    double[] data = img.get(y, x);
+                    ret.put(y, x, data[0], data[1], data[2], 255);
+                }
+            }
+        }
+        return new OpenCVImage(ret);
+    }
+
     /** {@inheritDoc} */
     @Override
     public Image getSubImage(int x, int y, int w, int h) {
@@ -185,6 +218,50 @@ class OpenCVImage implements Image {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void drawImage(Image overlay, boolean resize) {
+        if (!(overlay instanceof OpenCVImage)) {
+            throw new IllegalArgumentException("Only OpenCVImage allowed");
+        }
+        if (resize) {
+            overlay = overlay.resize(getWidth(), getHeight(), false);
+        }
+        Mat mat = (Mat) overlay.getWrappedImage();
+        if (mat.elemSize() != 4) {
+            mat.copyTo(image);
+            return;
+        }
+        int w = Math.min(image.width(), mat.width());
+        int h = Math.min(image.height(), mat.height());
+        for (int y = 0; y < h; ++y) {
+            for (int x = 0; x < w; ++x) {
+                /*
+                 * RA = SA + DA × (1 − SA)
+                 * R[0] = (S[0]×SA + D[0]×DA×(1 − SA)) / RA
+                 */
+                double[] src = mat.get(y, x);
+                double[] dest = image.get(y, x);
+                double sa = src[3];
+                double da;
+                double ra;
+                if (dest.length == 3) {
+                    da = 255 - sa;
+                    ra = 255;
+                } else {
+                    da = dest[3] * (255 - sa) / 255;
+                    ra = sa + da;
+                    dest[3] = ra;
+                }
+
+                dest[0] = (src[0] * sa + dest[0] * da) / ra;
+                dest[1] = (src[1] * sa + dest[1] * da) / ra;
+                dest[2] = (src[2] * sa + dest[2] * da) / ra;
+                image.put(y, x, dest);
+            }
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public List<BoundingBox> findBoundingBoxes() {
@@ -260,8 +337,21 @@ class OpenCVImage implements Image {
     private static Mat image2Mat(BufferedImage img) {
         int width = img.getWidth();
         int height = img.getHeight();
-        byte[] data = ((DataBufferByte) img.getRaster().getDataBuffer()).getData();
-        Mat mat = new Mat(height, width, CvType.CV_8UC3);
+        byte[] data;
+        Mat mat;
+        DataBuffer buf = img.getRaster().getDataBuffer();
+        if (buf instanceof DataBufferByte) {
+            data = ((DataBufferByte) buf).getData();
+            mat = new Mat(height, width, CvType.CV_8UC3);
+        } else if (buf instanceof DataBufferInt) {
+            int[] intData = ((DataBufferInt) buf).getData();
+            data = new byte[intData.length * 4];
+            ByteBuffer bb = ByteBuffer.wrap(data);
+            bb.asIntBuffer().put(intData);
+            mat = new Mat(height, width, CvType.CV_8UC4);
+        } else {
+            throw new IllegalArgumentException("Unsupported image type: " + buf.getClass());
+        }
         mat.put(0, 0, data);
         return mat;
     }
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
index 2b5baa3e9..9aef084ca 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
@@ -30,6 +30,8 @@ import org.opencv.imgproc.Imgproc;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.IntBuffer;
 import java.nio.file.Path;
 
 /** {@code OpenCVImageFactory} is a high performance implementation of {@link ImageFactory}. */
@@ -46,6 +48,7 @@ public class OpenCVImageFactory extends ImageFactory {
     /** {@inheritDoc} */
     @Override
     public Image fromFile(Path path) throws IOException {
+        // Load image without alpha channel
         Mat img = Imgcodecs.imread(path.toAbsolutePath().toString());
         if (img.empty()) {
             throw new IOException("Read image failed: " + path);
@@ -99,4 +102,20 @@ public class OpenCVImageFactory extends ImageFactory {
         Imgproc.cvtColor(img, img, Imgproc.COLOR_RGB2BGR);
         return new OpenCVImage(img);
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        Mat img = new Mat(height, width, CvType.CV_8UC4);
+        byte[] data = new byte[width * height * 4];
+        IntBuffer buf = ByteBuffer.wrap(data).asIntBuffer();
+        for (int pixel : pixels) {
+            int r = pixel >> 8 & 0xff00;
+            int g = pixel << 8 & 0xff0000;
+            int b = pixel << 24 & 0xff000000;
+            buf.put(pixel >>> 24 | b | g | r);
+        }
+        img.put(0, 0, data);
+        return new OpenCVImage(img);
+    }
 }


Test patch:
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
new file mode 100644
index 000000000..ec54e3bcb
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright 2022 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.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;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class SemanticSegmentationTranslatorFactoryTest {
+
+    private SemanticSegmentationTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new SemanticSegmentationTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, CategoryMask> translator =
+                    factory.newInstance(Image.class, CategoryMask.class, model, arguments);
+            Assert.assertTrue(translator instanceof SemanticSegmentationTranslator);
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
index 917fa088e..69c683bcf 100644
--- a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
+++ b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
@@ -25,9 +25,11 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.testing.TestRequirements;
 
+import org.opencv.core.Mat;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.awt.Color;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.file.Files;
@@ -156,4 +158,43 @@ public class OpenCVImageFactoryTest {
             }
         }
     }
+
+    @Test
+    public void testDrawImage() throws IOException {
+        ImageFactory factory = ImageFactory.getInstance();
+        int[] pixels = new int[64];
+        int index = 0;
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.RED.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.GREEN.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.BLUE.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.BLACK.getRGB() & 0x7fffffff;
+        }
+        Image img1 = factory.fromPixels(pixels, 8, 8);
+
+        pixels = new int[16];
+        for (int i = 0; i < 16; ++i) {
+            pixels[i] = Color.RED.getRGB() & 0x7fffffff;
+        }
+        Image img2 = factory.fromPixels(pixels, 2, 8);
+        Image mask = img2.getMask(new int[][] {{1, 1}, {1, 1}, {0, 0}, {0, 0}});
+        Image dup = img1.duplicate();
+        img1.drawImage(img2, false);
+        Mat mat = (Mat) img1.getWrappedImage();
+        byte[] data = new byte[64 * 4];
+        mat.get(0, 0, data);
+        Assert.assertEquals(data[2], -1); // red
+        Assert.assertEquals(data[3], -65); // alpha
+
+        dup.drawImage(mask, true);
+        ((Mat) dup.getWrappedImage()).get(0, 0, data);
+        Assert.assertEquals(data[2], -1); // red
+        Assert.assertEquals(data[3], -1); // alpha
+    }
 }