Panoptic Segmentation Model Export - EoMT¶

This notebook demonstrates how to export a panoptic segmentation model to ONNX and TensorRT.

The notebook covers the following steps:

Install LightlyTrain
Export a trained EoMT model to ONNX
Export a trained EoMT model to TensorRT
Run inference with the TensorRT engine

Important: When running on Google Colab make sure to select a GPU runtime. You can do this by going to Runtime > Change runtime type and selecting a GPU hardware accelerator.

!pip install "lightly-train[onnx,onnxruntime,onnxslim]"

Export to ONNX¶

Load the model weights¶

Then load the model with LightlyTrain’s load_model function. This will automatically download the model weights and load the model.

import lightly_train

model = lightly_train.load_model("dinov3/vits16-eomt-panoptic-coco")

Download an example image¶

Download an example image for inference with the following command:

!wget -O image.jpg http://images.cocodataset.org/val2017/000000105264.jpg

Preprocessing¶

import torch
import torchvision.transforms.v2 as T
from PIL import Image
from torchvision.transforms.functional import pil_to_tensor

# Load image with PIL.
image_pil = Image.open("image.jpg").convert("RGB")

# Convert PIL image to tensor for plotting.
image_tensor = pil_to_tensor(image_pil)

# Define pre-processing transforms.
w, h = image_pil.size
transforms = T.Compose(
    [
        T.Resize((model.image_size)),
        T.ToTensor(),
        T.Normalize(**model.image_normalize),
    ]
)

# Apply transforms for ONNX and TensorRT inference.
image_tensor_transformed = transforms(image_pil)[None]

Get the model predictions for reference¶

We define a helper function to visualize the predictions. The function will be used to compare the predictions from PyTorch, ONNX and TensorRT models.

import matplotlib.pyplot as plt
from torchvision.utils import draw_segmentation_masks


def visualize_segmentations(image, masks, segment_ids):
    # masks is (H, W, 2)
    # segment_ids is (num_segments)

    # Convert masks to boolean masks for each segment
    # masks[..., 1] contains the segment id for each pixel.
    # We add a mask for unassigned pixels (segment id -1).
    masks_bool = torch.stack(
        [masks[..., 1] == -1]
        + [masks[..., 1] == segment_id for segment_id in segment_ids]
    )

    # Create colors for visualization
    colors = [(0, 0, 0)] + [
        [int(color * 255) for color in plt.cm.tab20c(i / len(segment_ids))[:3]]
        for i in range(len(segment_ids))
    ]

    image_with_masks = draw_segmentation_masks(
        image, masks_bool, colors=colors, alpha=1.0
    )

    fig, axs = plt.subplots(1, 2, figsize=(12, 8))
    axs[0].imshow(image.permute(1, 2, 0))
    axs[0].axis("off")
    axs[1].imshow(image_with_masks.permute(1, 2, 0))
    axs[1].axis("off")
    fig.show()

# Get predictions from the PyTorch model.
results = model.predict(image_tensor)

# Visualize predictions from the PyTorch model.
visualize_segmentations(
    image_tensor, masks=results["masks"], segment_ids=results["segment_ids"]
)

Export the model to ONNX¶

# Export the PyTorch model to ONNX.
model.export_onnx(
    out="model.onnx",
    # precision="fp16", # Export model with FP16 weights for smaller size and faster inference.
)

See export_onnx for all available options when exporting to ONNX.

Run inference with the ONNX model¶

import onnxruntime as ort
import torch.nn.functional as F

# Create an ONNX Runtime session.
sess = ort.InferenceSession("model.onnx")

# Get expected input dtype.
input_dtype = sess.get_inputs()[0].type
input_dtype_numpy = {
    "tensor(float)": "float32",
    "tensor(float16)": "float16",
}[input_dtype]

# Run inference.
outputs = sess.run(
    output_names=None,
    input_feed={
        "images": image_tensor_transformed.numpy().astype(input_dtype_numpy),
    },
)

masks_onnx = outputs[0]
segment_ids_onnx = outputs[1]
scores_onnx = outputs[2]

# Resize ONNX predictions to original image size if necessary.
masks_onnx = torch.from_numpy(masks_onnx)
masks_onnx = masks_onnx.permute(0, 3, 1, 2).float()  # (1, 2, H_model, W_model)
# (1, 2, H_orig, W_orig)
masks_onnx = F.interpolate(masks_onnx, size=(h, w), mode="nearest")
masks_onnx = masks_onnx.permute(0, 2, 3, 1).long()  # (1, H_orig, W_orig, 2)

# Visualize predictions from the ONNX model.
visualize_segmentations(
    image_tensor, masks=masks_onnx[0], segment_ids=segment_ids_onnx[0]
)

Export to TensorRT¶

Requirements¶

TensorRT is not part of LightlyTrain’s dependencies and must be installed separately. Installation depends on your OS, Python version, GPU, and NVIDIA driver/CUDA setup. See the TensorRT documentation for more details.

On CUDA 12.x systems, install the TensorRT version tested with this tutorial. Pinning the version prevents pip from installing TensorRT 11, which is not yet supported by LightlyTrain:

!pip install "tensorrt-cu12==10.13.3.9"

# Get the TensorRT engine.
model.export_tensorrt(
    out="model.trt",
    # precision="fp16", # Export model with FP16 weights for smaller size and faster inference.
)

See export_tensorrt for all available options when exporting to TensorRT.

Run inference with the TensorRT engine¶

import numpy as np
import tensorrt as trt
import torch


class TRT:
    def __init__(
        self,
        engine_path: str,
        num_queries: int,
        device: str = "cuda:0",
        verbose: bool = False,
    ):
        self.device = torch.device(device)
        logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.INFO)
        trt.init_libnvinfer_plugins(logger, "")
        runtime = trt.Runtime(logger)

        with open(engine_path, "rb") as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

        io_names = [
            self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)
        ]
        self.in_names = [
            n
            for n in io_names
            if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT
        ]
        self.out_names = [
            n
            for n in io_names
            if self.engine.get_tensor_mode(n) == trt.TensorIOMode.OUTPUT
        ]

        self.buffers = {}
        self.bindings = []
        for name in io_names:
            shape = tuple(self.context.get_tensor_shape(name))
            shape = [s if s != -1 else num_queries for s in shape]
            np_dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            torch_dtype = torch.from_numpy(np.empty((), dtype=np_dtype)).dtype
            buffer = torch.empty(
                shape, device=self.device, dtype=torch_dtype
            ).contiguous()
            self.buffers[name] = buffer
            self.bindings.append(buffer.data_ptr())

    @torch.no_grad()
    def __call__(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        for name in self.in_names:
            self.buffers[name].copy_(inputs[name].to(self.device))
        if not self.context.execute_v2(self.bindings):
            raise RuntimeError("TensorRT execution failed")
        return {name: self.buffers[name] for name in self.out_names}

# Instantiate the TensorRT model.
trt_model = TRT("model.trt", num_queries=model.num_queries)

# Run inference with the TensorRT model.
outputs_trt = trt_model({"images": image_tensor_transformed})

masks_trt = outputs_trt["masks"][0]
segment_ids_trt = outputs_trt["segment_ids"][0]
scores_trt = outputs_trt["scores"][0]

# Filter out segments with scores below threshold. This is required because TensorRT
# returns scores and segment ids for all possible segments.
keep = scores_trt > 0.8
segment_ids_trt = segment_ids_trt[keep]
scores_trt = scores_trt[keep]

# Resize TensorRT predictions to original image size if necessary.
masks_trt = masks_trt.permute(2, 0, 1).float()
masks_trt = masks_trt.unsqueeze(0)
masks_trt = F.interpolate(masks_trt, size=(h, w), mode="nearest")
masks_trt = masks_trt.squeeze(0).permute(1, 2, 0).long()

# Visualize predictions from the TensorRT model.
visualize_segmentations(
    image_tensor, masks=masks_trt.cpu(), segment_ids=segment_ids_trt.cpu()
)