Prepare Segment Anything ONNX Model

There are two Segment Anything models available to select from in the steps below: SAM and Mobile SAM. SAM provides the full accuracy of Segment Anything but can consume over 8GB of GPU memory and take more computation effort for inference. Mobile SAM, has been tuned to operate with a much smaller memory footprint and less computation overhead, but with some level of degradation in quality. Choose the model which will be the most effective for your use case.

1. Make a directory to place models (inside the Isaac ROS environment):

   mkdir -p ${ISAAC_ROS_WS}/isaac_ros_assets/models/triton/segment_anything/1/
   

2. Create the ONNX file for SAM or for Mobile SAM.

Note

SAM is only supported with Triton using the ONNX backend. The following step is to be performed only on x86_64. You must copy the generated ONNX file to a Jetson, in case you want to run the package on the device.

Note

It is recommended to run SAM with at least 12GB+ GPU memory. Mobile SAM smaller version of SAM which can be used if GPU VRAM is not sufficient to run SAM.

SAMMobileSAM

Steps to use SAM

To create ONNX file for SAM. Download the PyTorch weights from Segment Anything official repository. Then, copy the downloaded model weights into the container. The model is expected to be downloaded to ~/Downloads outside the Isaac ROS environment.

Move this file to ${ISAAC_ROS_WS}/isaac_ros_assets/isaac_ros_segment_anything/ using a console outside of the container and then continue inside the container.

For example, if the model vit_b.pth was downloaded to ~/Downloads:

On x86_64:

# Outside container
mv ~/Downloads/</path/to/vit_b> ${ISAAC_ROS_WS}/isaac_ros_assets/isaac_ros_segment_anything/vit_b.pth

# Inside container
cd ${ISAAC_ROS_WS}
pip3 install --break-system-packages git+https://github.com/facebookresearch/segment-anything.git && \
pip3 install --break-system-packages torchvision

# Inside container
ros2 run isaac_ros_segment_anything torch_to_onnx.py \
--checkpoint ${ISAAC_ROS_WS}/isaac_ros_assets/isaac_ros_segment_anything/vit_b.pth \
--output ${ISAAC_ROS_WS}/isaac_ros_assets/models/triton/segment_anything/1/model.onnx --model-type vit_b --sam-type SAM

1. Copy the config file:

   cp ${ISAAC_ROS_WS}/isaac_ros_assets/isaac_ros_segment_anything/sam_config_onnx.pbtxt ${ISAAC_ROS_WS}/isaac_ros_assets/models/triton/segment_anything/config.pbtxt