Repository: deepjavalibrary/djl
Pull request #3495 (2024-10-04 00:32:44)

[api] Improve Sam2Translator for PyTorch traced model
- **[api] Adds sam2 model to onnxruntime model zoo**
- **[api] Improve Sam2Translator for PyTorch traced model**

## 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



Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/Sam2ServingTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/Sam2ServingTranslator.java
index 909c3d01f..97432dd03 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/Sam2ServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/Sam2ServingTranslator.java
@@ -14,7 +14,9 @@ package ai.djl.modality.cv.translator;
 
 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.TranslateException;
 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 class Sam2ServingTranslator implements Translator<Input, Output> {
 
     /** {@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 class Sam2ServingTranslator implements Translator<Input, Output> {
                 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);
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 0c05c569f..751019c76 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
@@ -58,6 +58,7 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
     private Pipeline pipeline;
     private Predictor<NDList, NDList> predictor;
     private String encoderPath;
+    private String encodeMethod;
 
     /** Constructs a {@code Sam2Translator} instance. */
     public Sam2Translator(Builder builder) {
@@ -66,12 +67,19 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
         pipeline.add(new ToTensor());
         pipeline.add(new Normalize(MEAN, STD));
         this.encoderPath = builder.encoderPath;
+        this.encodeMethod = builder.encodeMethod;
     }
 
     /** {@inheritDoc} */
     @Override
     public void prepare(TranslatorContext ctx) throws IOException, ModelException {
         if (encoderPath == null) {
+            // PyTorch model
+            if (encodeMethod != null) {
+                Model model = ctx.getModel();
+                predictor = model.newPredictor(new NoopTranslator(null));
+                model.getNDManager().attachInternal(UUID.randomUUID().toString(), predictor);
+            }
             return;
         }
         Model model = ctx.getModel();
@@ -111,7 +119,15 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             return new NDList(array, locations, labels);
         }
 
-        NDList embeddings = predictor.predict(new NDList(array));
+        NDList embeddings;
+        if (encodeMethod == null) {
+            embeddings = predictor.predict(new NDList(array));
+        } else {
+            NDArray placeholder = manager.create("");
+            placeholder.setName("module_method:" + encodeMethod);
+            embeddings = predictor.predict(new NDList(placeholder, array));
+        }
+
         NDArray mask = manager.zeros(new Shape(1, 1, 256, 256));
         NDArray hasMask = manager.zeros(new Shape(1));
         return new NDList(
@@ -173,9 +189,11 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
     public static class Builder {
 
         String encoderPath;
+        String encodeMethod;
 
         Builder(Map<String, ?> arguments) {
             encoderPath = ArgumentsUtil.stringValue(arguments, "encoder");
+            encodeMethod = ArgumentsUtil.stringValue(arguments, "encode_method");
         }
 
         /**
@@ -189,6 +207,17 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             return this;
         }
 
+        /**
+         * Sets the module name for encode method.
+         *
+         * @param encodeMethod the module name for encode method
+         * @return the builder
+         */
+        public Builder optEncodeMethod(String encodeMethod) {
+            this.encodeMethod = encodeMethod;
+            return this;
+        }
+
         /**
          * Builds the translator.
          *
@@ -205,6 +234,7 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
         private Image image;
         private Point[] points;
         private int[] labels;
+        private boolean visualize;
 
         /**
          * Constructs a {@code Sam2Input} instance.
@@ -214,9 +244,22 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
          * @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 +271,15 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             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 +340,16 @@ 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 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 +377,7 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             private Image image;
             private List<Point> points;
             private List<Integer> labels;
+            private boolean visualize;
 
             Builder(Image image) {
                 this.image = image;
@@ -380,6 +436,16 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
                 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 +454,7 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             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 +479,7 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
         private static final class Prompt {
             String image;
             Location[] prompt;
+            boolean visualize;
 
             public void setImage(String image) {
                 this.image = image;
@@ -421,6 +488,10 @@ public class Sam2Translator implements NoBatchifyTranslator<Sam2Input, DetectedO
             public void setPrompt(Location[] prompt) {
                 this.prompt = prompt;
             }
+
+            public void setVisualize(boolean visualize) {
+                this.visualize = visualize;
+            }
         }
     }
 }
diff --git a/api/src/main/java/ai/djl/util/JsonSerializable.java b/api/src/main/java/ai/djl/util/JsonSerializable.java
index b998e94e5..e321ec21e 100644
--- a/api/src/main/java/ai/djl/util/JsonSerializable.java
+++ b/api/src/main/java/ai/djl/util/JsonSerializable.java
@@ -49,11 +49,7 @@ public interface JsonSerializable extends Serializable, BytesSupplier {
         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. */
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
index b92fea429..5ea9920b8 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
@@ -43,10 +43,8 @@ public class PtModelZoo extends ModelZoo {
         addModel(
                 REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet18_embedding", "0.0.1"));
         addModel(REPOSITORY.model(CV.INSTANCE_SEGMENTATION, GROUP_ID, "yolov8n-seg", "0.0.1"));
-        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-tiny", "0.0.1"));
-        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-tiny-gpu", "0.0.1"));
-        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-large", "0.0.1"));
-        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-large-gpu", "0.0.1"));
+        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-tiny", "0.0.2"));
+        addModel(REPOSITORY.model(CV.MASK_GENERATION, GROUP_ID, "sam2-hiera-large", "0.0.2"));
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1"));
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolov5s", "0.0.1"));
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolov8n", "0.0.1"));
diff --git a/examples/docs/segment_anything_2.md b/examples/docs/segment_anything_2.md
index 768a666b8..6c66bbe91 100644
--- a/examples/docs/segment_anything_2.md
+++ b/examples/docs/segment_anything_2.md
@@ -73,15 +73,21 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 from torch import nn
 
 
-class SAM2ImageEncoder(nn.Module):
+class Sam2Wrapper(nn.Module):
 
-    def __init__(self, sam_model: SAM2Base) -> None:
+    def __init__(self, sam_model: SAM2Base, multimask_output: bool) -> None:
         super().__init__()
         self.model = sam_model
         self.image_encoder = sam_model.image_encoder
         self.no_mem_embed = sam_model.no_mem_embed
+        self.mask_decoder = sam_model.sam_mask_decoder
+        self.prompt_encoder = sam_model.sam_prompt_encoder
+        self.img_size = sam_model.image_size
+        self.multimask_output = multimask_output
+        self.sparse_embedding = None
 
-    def forward(self, x: torch.Tensor) -> tuple[Any, Any, Any]:
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> tuple[Any, Any, Any]:
         backbone_out = self.image_encoder(x)
         backbone_out["backbone_fpn"][0] = self.model.sam_mask_decoder.conv_s0(
             backbone_out["backbone_fpn"][0])
@@ -106,18 +112,6 @@ class SAM2ImageEncoder(nn.Module):
 
         return feats[0], feats[1], feats[2]
 
-
-class SAM2ImageDecoder(nn.Module):
-
-    def __init__(self, sam_model: SAM2Base, multimask_output: bool) -> None:
-        super().__init__()
-        self.mask_decoder = sam_model.sam_mask_decoder
-        self.prompt_encoder = sam_model.sam_prompt_encoder
-        self.model = sam_model
-        self.img_size = sam_model.image_size
-        self.multimask_output = multimask_output
-        self.sparse_embedding = None
-
     @torch.no_grad()
     def forward(
         self,
@@ -205,17 +199,13 @@ def trace_model(model_id: str):
         device = torch.device("cpu")
 
     model_name = f"{model_id[9:]}"
-    os.makedirs(model_name)
+    os.makedirs(model_name, exist_ok=True)
 
     predictor = SAM2ImagePredictor.from_pretrained(model_id, device=device)
-    encoder = SAM2ImageEncoder(predictor.model)
-    decoder = SAM2ImageDecoder(predictor.model, True)
+    model = Sam2Wrapper(predictor.model, True)
 
     input_image = torch.ones(1, 3, 1024, 1024).to(device)
-    high_res_feats_0, high_res_feats_1, image_embed = encoder(input_image)
-
-    converted = torch.jit.trace(encoder, input_image)
-    torch.jit.save(converted, f"model_name/encoder.pt")
+    high_res_feats_0, high_res_feats_1, image_embed = model.encode(input_image)
 
     # trace decoder model
     embed_size = (
@@ -232,10 +222,14 @@ def trace_model(model_id: str):
     mask_input = torch.randn(1, 1, *mask_input_size, dtype=torch.float)
     has_mask_input = torch.tensor([1], dtype=torch.float)
 
-    converted = torch.jit.trace(
-        decoder, (image_embed, high_res_feats_0, high_res_feats_1,
-                  point_coords, point_labels, mask_input, has_mask_input))
-    torch.jit.save(converted, f"model_name/model_name.pt")
+    converted = torch.jit.trace_module(
+        model, {
+            "encode":
+            input_image,
+            "forward": (image_embed, high_res_feats_0, high_res_feats_1,
+                        point_coords, point_labels, mask_input, has_mask_input)
+        })
+    torch.jit.save(converted, f"{model_name}/{model_name}.pt")
 
     # save serving.properties
     serving_file = os.path.join(model_name, "serving.properties")
@@ -244,7 +238,7 @@ def trace_model(model_id: str):
             f"engine=PyTorch\n"
             f"option.modelName={model_name}\n"
             f"translatorFactory=ai.djl.modality.cv.translator.Sam2TranslatorFactory\n"
-            f"encoder=encoder.pt")
+            f"encode_method=encode\n")
 
 
 if __name__ == '__main__':
diff --git a/examples/docs/trace_sam2_img.py b/examples/docs/trace_sam2_img.py
index ce5a94ed0..650284cc1 100644
--- a/examples/docs/trace_sam2_img.py
+++ b/examples/docs/trace_sam2_img.py
@@ -20,15 +20,21 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 from torch import nn
 
 
-class SAM2ImageEncoder(nn.Module):
+class Sam2Wrapper(nn.Module):
 
-    def __init__(self, sam_model: SAM2Base) -> None:
+    def __init__(self, sam_model: SAM2Base, multimask_output: bool) -> None:
         super().__init__()
         self.model = sam_model
         self.image_encoder = sam_model.image_encoder
         self.no_mem_embed = sam_model.no_mem_embed
+        self.mask_decoder = sam_model.sam_mask_decoder
+        self.prompt_encoder = sam_model.sam_prompt_encoder
+        self.img_size = sam_model.image_size
+        self.multimask_output = multimask_output
+        self.sparse_embedding = None
 
-    def forward(self, x: torch.Tensor) -> tuple[Any, Any, Any]:
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> tuple[Any, Any, Any]:
         backbone_out = self.image_encoder(x)
         backbone_out["backbone_fpn"][0] = self.model.sam_mask_decoder.conv_s0(
             backbone_out["backbone_fpn"][0])
@@ -53,18 +59,6 @@ class SAM2ImageEncoder(nn.Module):
 
         return feats[0], feats[1], feats[2]
 
-
-class SAM2ImageDecoder(nn.Module):
-
-    def __init__(self, sam_model: SAM2Base, multimask_output: bool) -> None:
-        super().__init__()
-        self.mask_decoder = sam_model.sam_mask_decoder
-        self.prompt_encoder = sam_model.sam_prompt_encoder
-        self.model = sam_model
-        self.img_size = sam_model.image_size
-        self.multimask_output = multimask_output
-        self.sparse_embedding = None
-
     @torch.no_grad()
     def forward(
         self,
@@ -152,17 +146,13 @@ def trace_model(model_id: str):
         device = torch.device("cpu")
 
     model_name = f"{model_id[9:]}"
-    os.makedirs(model_name)
+    os.makedirs(model_name, exist_ok=True)
 
     predictor = SAM2ImagePredictor.from_pretrained(model_id, device=device)
-    encoder = SAM2ImageEncoder(predictor.model)
-    decoder = SAM2ImageDecoder(predictor.model, True)
+    model = Sam2Wrapper(predictor.model, True)
 
     input_image = torch.ones(1, 3, 1024, 1024).to(device)
-    high_res_feats_0, high_res_feats_1, image_embed = encoder(input_image)
-
-    converted = torch.jit.trace(encoder, input_image)
-    torch.jit.save(converted, f"model_name/encoder.pt")
+    high_res_feats_0, high_res_feats_1, image_embed = model.encode(input_image)
 
     # trace decoder model
     embed_size = (
@@ -179,10 +169,14 @@ def trace_model(model_id: str):
     mask_input = torch.randn(1, 1, *mask_input_size, dtype=torch.float)
     has_mask_input = torch.tensor([1], dtype=torch.float)
 
-    converted = torch.jit.trace(
-        decoder, (image_embed, high_res_feats_0, high_res_feats_1,
-                  point_coords, point_labels, mask_input, has_mask_input))
-    torch.jit.save(converted, f"model_name/model_name.pt")
+    converted = torch.jit.trace_module(
+        model, {
+            "encode":
+            input_image,
+            "forward": (image_embed, high_res_feats_0, high_res_feats_1,
+                        point_coords, point_labels, mask_input, has_mask_input)
+        })
+    torch.jit.save(converted, f"{model_name}/{model_name}.pt")
 
     # save serving.properties
     serving_file = os.path.join(model_name, "serving.properties")
@@ -191,7 +185,7 @@ def trace_model(model_id: str):
             f"engine=PyTorch\n"
             f"option.modelName={model_name}\n"
             f"translatorFactory=ai.djl.modality.cv.translator.Sam2TranslatorFactory\n"
-            f"encoder=encoder.pt")
+            f"encode_method=encode\n")
 
 
 if __name__ == '__main__':


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 9f9811bda..a742e92d4 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
@@ -33,15 +33,17 @@ public class Sam2InputTest {
                 "{\"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);
     }