[api] Visualize sam2 output for Sam2ServingTranslator

api/src/main/java/ai/djl/modality/cv/translator/Sam2ServingTranslator.java

@@ -14,7 +14,9 @@
 
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.Rectangle;
 import ai.djl.modality.cv.translator.Sam2Translator.Sam2Input;
 import ai.djl.ndarray.BytesSupplier;
 import ai.djl.ndarray.NDList;
@@ -23,7 +25,13 @@
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
 
+import org.apache.commons.codec.binary.Base64OutputStream;
+
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
 /** A {@link Translator} that can serve SAM2 model. */
 public class Sam2ServingTranslator implements Translator<Input, Output> {
@@ -47,11 +55,29 @@ public Batchifier getBatchifier() {
 
     /** {@inheritDoc} */
     @Override
-    public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
+    public Output processOutput(TranslatorContext ctx, NDList list) throws IOException {
         Output output = new Output();
+        Sam2Input sam2 = (Sam2Input) ctx.getAttachment("input");
         output.addProperty("Content-Type", "application/json");
-        DetectedObjects obj = translator.processOutput(ctx, list);
-        output.add(BytesSupplier.wrapAsJson(obj));
+        DetectedObjects detection = translator.processOutput(ctx, list);
+        Map<String, Object> ret = new LinkedHashMap<>(); // NOPMD
+        ret.put("result", detection);
+        if (sam2.isVisualize()) {
+            Image img = sam2.getImage();
+            img.drawBoundingBoxes(detection, 0.8f);
+            img.drawMarks(sam2.getPoints());
+            for (Rectangle rect : sam2.getBoxes()) {
+                img.drawRectangle(rect, 0xff0000, 6);
+            }
+            ByteArrayOutputStream os = new ByteArrayOutputStream();
+            os.write("data:image/png;base64,".getBytes(StandardCharsets.UTF_8));
+            Base64OutputStream bos = new Base64OutputStream(os, true, 0, null);
+            img.save(bos, "png");
+            bos.close();
+            os.close();
+            ret.put("image", os.toString(StandardCharsets.UTF_8.name()));
+        }
+        output.add(BytesSupplier.wrapAsJson(ret));
         return output;
     }
 
@@ -64,6 +90,7 @@ public NDList processInput(TranslatorContext ctx, Input input) throws Exception
                 throw new TranslateException("Input data is empty.");
             }
             Sam2Input sam2 = Sam2Input.fromJson(data.getAsString());
+            ctx.setAttachment("input", sam2);
             return translator.processInput(ctx, sam2);
         } catch (IOException e) {
             throw new TranslateException("Input is not an Image data type", e);

api/src/main/java/ai/djl/modality/cv/translator/Sam2Translator.java

@@ -205,6 +205,7 @@ public static final class Sam2Input {
         private Image image;
         private Point[] points;
         private int[] labels;
+        private boolean visualize;
 
         /**
          * Constructs a {@code Sam2Input} instance.
@@ -214,9 +215,22 @@ public static final class Sam2Input {
          * @param labels the labels for the locations (0: background, 1: foreground)
          */
         public Sam2Input(Image image, Point[] points, int[] labels) {
+            this(image, points, labels, false);
+        }
+
+        /**
+         * Constructs a {@code Sam2Input} instance.
+         *
+         * @param image the image
+         * @param points the locations on the image
+         * @param labels the labels for the locations (0: background, 1: foreground)
+         * @param visualize true if output visualized image
+         */
+        public Sam2Input(Image image, Point[] points, int[] labels, boolean visualize) {
             this.image = image;
             this.points = points;
             this.labels = labels;
+            this.visualize = visualize;
         }
 
         /**
@@ -228,6 +242,15 @@ public Image getImage() {
             return image;
         }
 
+        /**
+         * Returns {@code true} if output visualized image.
+         *
+         * @return {@code true} if output visualized image
+         */
+        public boolean isVisualize() {
+            return visualize;
+        }
+
         /**
          * Returns the locations.
          *
@@ -288,13 +311,16 @@ 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 url value");
+                throw new IllegalArgumentException("Missing image value");
             }
             if (prompt.prompt == null || prompt.prompt.length == 0) {
                 throw new IllegalArgumentException("Missing prompt value");
             }
             Image image = ImageFactory.getInstance().fromUrl(prompt.image);
             Builder builder = builder(image);
+            if (prompt.visualize) {
+                builder.visualize();
+            }
             for (Location location : prompt.prompt) {
                 int[] data = location.data;
                 if ("point".equals(location.type)) {
@@ -322,6 +348,7 @@ public static final class Builder {
             private Image image;
             private List<Point> points;
             private List<Integer> labels;
+            private boolean visualize;
 
             Builder(Image image) {
                 this.image = image;
@@ -380,6 +407,16 @@ public Builder addBox(int x, int y, int right, int bottom) {
                 return this;
             }
 
+            /**
+             * Sets the visualize for the {@code Sam2Input}.
+             *
+             * @return the builder
+             */
+            public Builder visualize() {
+                visualize = true;
+                return this;
+            }
+
             /**
              * Builds the {@code Sam2Input}.
              *
@@ -388,7 +425,7 @@ public Builder addBox(int x, int y, int right, int bottom) {
             public Sam2Input build() {
                 Point[] location = points.toArray(new Point[0]);
                 int[] array = labels.stream().mapToInt(Integer::intValue).toArray();
-                return new Sam2Input(image, location, array);
+                return new Sam2Input(image, location, array, visualize);
             }
         }
 
@@ -413,6 +450,7 @@ public void setLabel(int label) {
         private static final class Prompt {
             String image;
             Location[] prompt;
+            boolean visualize;
 
             public void setImage(String image) {
                 this.image = image;
@@ -421,6 +459,10 @@ public void setImage(String image) {
             public void setPrompt(Location[] prompt) {
                 this.prompt = prompt;
             }
+
+            public void setVisualize(boolean visualize) {
+                this.visualize = visualize;
+            }
         }
     }
 }

api/src/main/java/ai/djl/util/JsonSerializable.java

@@ -49,11 +49,7 @@ default ByteBuffer toByteBuffer() {
         return ByteBuffer.wrap(toJson().getBytes(StandardCharsets.UTF_8));
     }
 
-    /**
-     * Serializes the object to the {@code JsonElement}.
-     *
-     * @return the {@code JsonElement}
-     */
+    /** {@inheritDoc} */
     JsonElement serialize();
 
     /** A customized Gson serializer to serialize the {@code Segmentation} object. */

api/src/test/java/ai/djl/modality/cv/translator/Sam2InputTest.java

@@ -33,15 +33,17 @@ public void test() throws IOException {
                 "{\"image\": \""
                         + file.toUri().toURL()
                         + "\",\n"
+                        + "\"visualize\": true,\n"
                         + "\"prompt\": [\n"
                         + "    {\"type\": \"point\", \"data\": [575, 750], \"label\": 0},\n"
                         + "    {\"type\": \"rectangle\", \"data\": [425, 600, 700, 875]}\n"
                         + "]}";
         Sam2Input input = Sam2Input.fromJson(json);
+        Assert.assertTrue(input.isVisualize());
         Assert.assertEquals(input.getPoints().size(), 1);
         Assert.assertEquals(input.getBoxes().size(), 1);
 
-        input = Sam2Input.builder(img).addPoint(0, 1).addBox(0, 0, 1, 1).build();
+        input = Sam2Input.builder(img).visualize().addPoint(0, 1).addBox(0, 0, 1, 1).build();
         Assert.assertEquals(input.getPoints().size(), 1);
         Assert.assertEquals(input.getBoxes().size(), 1);
     }