Add SAM OVMS and Triton server Docker image builders

.github/codecov.yml

@@ -0,0 +1,8 @@
+coverage:
+  status:
+    project:
+      default:
+        target: auto
+        # adjust accordingly based on how flaky your tests are
+        # this allows a 2% drop from the previous base commit coverage
+        threshold: 2%

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (<https://github.com/openvinotoolkit/datumaro/pull/1089>)
 - Support video annotation import/export
   (<https://github.com/openvinotoolkit/datumaro/pull/1124>)
+- Add SAM OVMS and Triton server Docker image builders
+  (<https://github.com/openvinotoolkit/datumaro/pull/1129>)
 
 ### Enhancements
 - Remove xfail marks from the convert integration tests

docker/segment-anything/Dockerfile.exporter

@@ -0,0 +1,44 @@
+ARG model_type
+
+FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime AS onnx-exporter-base
+
+ARG http_proxy
+ARG https_proxy
+ARG no_proxy
+
+RUN apt update && apt install -y git curl && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN pip install opencv-python pycocotools matplotlib onnxruntime onnx
+
+WORKDIR /ws
+RUN mkdir -p /ws/output
+RUN git clone https://github.com/facebookresearch/segment-anything.git
+
+WORKDIR /ws/segment-anything
+RUN pip install -e .
+
+# From https://github.com/facebookresearch/segment-anything/tree/main#model-checkpoints
+# vit_h: https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
+# vit_l: https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth
+# vit_b: https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
+
+FROM onnx-exporter-base AS onnx-exporter-vit_h
+
+RUN curl -o ckpt.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
+
+FROM onnx-exporter-base AS onnx-exporter-vit_l
+
+RUN curl -o ckpt.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth
+
+FROM onnx-exporter-base AS onnx-exporter-vit_b
+
+RUN curl -o ckpt.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
+
+FROM onnx-exporter-${model_type} AS onnx-exporter-final
+ARG model_type
+
+RUN python scripts/export_onnx_model.py --checkpoint ckpt.pth --model-type ${model_type} --return-single-mask --output /ws/output/decoder.onnx
+
+COPY export_onnx_encoder.py scripts/
+RUN python scripts/export_onnx_encoder.py --checkpoint ckpt.pth --model-type ${model_type} --output /ws/output/encoder.onnx

docker/segment-anything/Dockerfile.ovms

@@ -0,0 +1,15 @@
+FROM openvino/model_server:2023.0-gpu AS ovms
+
+USER root
+
+RUN mkdir -p /models/sam_encoder/1
+RUN mkdir -p /models/sam_decoder/1
+
+COPY --from=onnx-exporter /ws/output/encoder.onnx /models/sam_encoder/1/model.onnx
+COPY --from=onnx-exporter /ws/output/decoder.onnx /models/sam_decoder/1/model.onnx
+
+ADD ovms/model_config_list.json /models/model_config_list.json
+
+USER ovms
+
+ENTRYPOINT ["/ovms/bin/ovms", "--config_path", "/models/model_config_list.json"]

docker/segment-anything/Dockerfile.triton

@@ -0,0 +1,12 @@
+FROM nvcr.io/nvidia/tritonserver:23.07-py3 AS nvidia-triton
+
+ADD triton/sam_encoder /models/sam_encoder
+ADD triton/sam_decoder /models/sam_decoder
+
+RUN mkdir -p /models/sam_encoder/1
+RUN mkdir -p /models/sam_decoder/1
+
+COPY --from=onnx-exporter /ws/output/encoder.onnx /models/sam_encoder/1/model.onnx
+COPY --from=onnx-exporter /ws/output/decoder.onnx /models/sam_decoder/1/model.onnx
+
+ENTRYPOINT [ "tritonserver", "--model-repository=/models" ]

docker/segment-anything/README.md

@@ -0,0 +1,78 @@
+# Segment Anything Model for Datumaro
+
+## Introduction
+
+[Segment Anything Model](https://github.com/facebookresearch/segment-anything) provides high quality object masks for the given point, bounding box, and text prompts.
+[Datumaro](https://github.com/openvinotoolkit/datumaro) utilzes this zero-shot generalizable model for the dataset transformation features such as:
+
+1. Bounding box to instance segmentation transform
+2. Pseudo labeling for instance segmentation
+3. ...
+
+## Build Docker image
+
+### Prerequisite
+
+1. You should [install Docker](https://docs.docker.com/engine/install/ubuntu/) in your machine.
+    In addition, we recommend you to use [Ubuntu](https://ubuntu.com/) since we only provide Linux shell script for the Docker image building.
+
+2. Clone Datumaro repository to your local:
+    ```console
+    git clone https://github.com/openvinotoolkit/datumaro
+    ```
+
+3. go to `docker/segment-anything` sub-directory
+    ```console
+    cd docker/segment-anything
+    ```
+
+### Building Docker image for OpenVINO™ Model Server
+
+1. Execute `build_ovms.sh` shell script with `MODEL_TYPE` argument:
+    ```console
+    # MODEL_TYPE := "vit_h", "vit_l", or "vit_b"
+
+    ./build_ovms.sh <MODEL_TYPE>
+    ```
+2. It will create a Docker image to your local repository, which is named as `segment-anything-ovms:<MODEL_TYPE>`.
+
+### Building Docker image for NVIDIA Triton™ Inference Server
+
+1. Execute `build_ovms.sh` shell script with `MODEL_TYPE` argument:
+    ```console
+    # MODEL_TYPE := "vit_h", "vit_l", or "vit_b"
+
+    ./build_triton.sh <MODEL_TYPE>
+    ```
+2. It will create a Docker image to your local repository, which is named as `segment-anything-triton-server:<MODEL_TYPE>`.
+
+## Launch model server instances
+
+We provide how to launch the dedicated model server instance for Segment Anything Model.
+Because we built the Docker image for it, you can scale out it with the container orchestration system such as Kubernetes.
+However, in this guide, we help you how to launch a single model server instance in your local machine.
+Nevertheless, if you have multiple accelerators in your machine (e.g., 8 GPUs for a machine),
+OpenVINO™ Model Server or NVIDIA Triton™ Inference Server either can use all installed devices.
+For more details, please see the guide of these model server solutions.
+
+### Building Docker image for OpenVINO™ Model Server
+
+Execute the following command in your shell to launch the inference server:
+
+```console
+# MODEL_TYPE := "vit_h", "vit_l", or "vit_b"
+docker run -it --rm -p 9000:9000 segment-anything-ovms:<MODEL_TYPE> --port 9000
+```
+
+### Building Docker image for NVIDIA Triton™ Inference Server
+
+Execute the following command in your shell to launch the inference server:
+
+```console
+# MODEL_TYPE := "vit_h", "vit_l", or "vit_b"
+docker run --gpus=all -it --shm-size=256m --rm -p8000:8000 -p8001:8001 -p8002:8002 segment-anything-triton-server:<MODEL_TYPE>
+```
+
+## Can I use the SAM model server I built standalone?
+
+Yes, you can use or test your SAM model server instances. We provide a detailed examples for it in [this notebook example](./test_model_servers.ipynb).

docker/segment-anything/build_ovms.sh

@@ -0,0 +1,19 @@
+if [ -z "$1" ] ; then
+    echo "No model type is given (availables: vit_h, vit_l, vit_b)"
+    exit 1
+fi
+
+MODEL_TYPE=$1
+
+echo "Build MODEL_TYPE=$1"
+
+docker build \
+    --build-arg http_proxy=${http_proxy} \
+    --build-arg https_proxy=${https_proxy} \
+    --build-arg no_proxy=${no_proxy} \
+    --build-arg model_type=${MODEL_TYPE} \
+    -f Dockerfile.exporter -t segment-anything-onnx-exporter:${MODEL_TYPE} .
+
+docker build \
+    --build-context onnx-exporter=docker-image://segment-anything-onnx-exporter:${MODEL_TYPE} \
+    -f Dockerfile.ovms -t segment-anything-ovms:${MODEL_TYPE} .

docker/segment-anything/build_triton.sh

@@ -0,0 +1,19 @@
+if [ -z "$1" ] ; then
+    echo "No model type is given (availables: vit_h, vit_l, vit_b)"
+    exit 1
+fi
+
+MODEL_TYPE=$1
+
+echo "Build MODEL_TYPE=$1"
+
+docker build \
+    --build-arg http_proxy=${http_proxy} \
+    --build-arg https_proxy=${https_proxy} \
+    --build-arg no_proxy=${no_proxy} \
+    --build-arg model_type=${MODEL_TYPE} \
+    -f Dockerfile.exporter -t segment-anything-onnx-exporter:${MODEL_TYPE} .
+
+docker build \
+    --build-context onnx-exporter=docker-image://segment-anything-onnx-exporter:${MODEL_TYPE} \
+    -f Dockerfile.triton -t segment-anything-triton-server:${MODEL_TYPE} .

docker/segment-anything/export_onnx_encoder.py

@@ -0,0 +1,72 @@
+# Copyright (C) 2023 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import argparse
+
+import numpy as np
+import torch
+from segment_anything import sam_model_registry
+from torch import nn
+
+parser = argparse.ArgumentParser(description="Export the SAM image encoder to an ONNX model.")
+
+parser.add_argument(
+    "--checkpoint",
+    type=str,
+    required=True,
+    help="The path to the SAM model checkpoint.",
+)
+
+parser.add_argument(
+    "--output", type=str, required=True, help="The filename to save the ONNX model to."
+)
+
+parser.add_argument(
+    "--model-type",
+    type=str,
+    required=True,
+    help="In ['default', 'vit_h', 'vit_l', 'vit_b']. Which type of SAM model to export.",
+)
+
+parser.add_argument(
+    "--opset",
+    type=int,
+    default=17,
+    help="The ONNX opset version to use. Must be >=11",
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    class Encoder(nn.Module):
+        def __init__(self, model_type: str, checkpoint: str) -> None:
+            super().__init__()
+            self.sam = sam_model_registry[model_type](checkpoint=checkpoint)
+
+        def forward(self, img: torch.Tensor):
+            x = self.sam.preprocess(img)
+            return self.sam.image_encoder(x)
+
+    inputs = {
+        "img": torch.from_numpy(
+            np.random.randint(0, 256, size=[1, 3, 1024, 1024], dtype=np.uint8)
+        ).to(dtype=torch.float32)
+    }
+
+    output_names = ["image_embeddings"]
+
+    with open(args.output, "wb") as f:
+        torch.onnx.export(
+            Encoder(model_type=args.model_type, checkpoint=args.checkpoint),
+            inputs,
+            f,
+            export_params=True,
+            verbose=False,
+            opset_version=args.opset,
+            do_constant_folding=True,
+            input_names=list(inputs.keys()),
+            output_names=output_names,
+            dynamic_axes={"img": [2, 3]},
+        )

docker/segment-anything/ovms/model_config_list.json

@@ -0,0 +1,23 @@
+{
+    "model_config_list": [
+        {
+            "config": {
+                "name": "sam_encoder",
+                "base_path": "/models/sam_encoder",
+                "layout": {
+                    "img": "NCHW"
+                }
+            }
+        },
+        {
+            "config": {
+                "name": "sam_decoder",
+                "base_path": "/models/sam_decoder",
+                "layout": {
+                    "has_mask_input": "...",
+                    "orig_im_size": "..."
+                }
+            }
+        }
+    ]
+}