首先导出onnx模型，结构如下：

onnxruntime推理

image_preprocess.py

import numpy as np
import cv2

def preprocess_image(image):
    """Preprocesses the image for the model."""
    image = cv2.resize(image, (1008, 1008), interpolation=cv2.INTER_LINEAR)

    image = image.astype(np.float32) / 255.0
    image = (image - np.array([0.5, 0.5, 0.5])) / np.array([0.5, 0.5, 0.5])

    # 从 HWC → CHW
    image = np.transpose(image, (2, 0, 1))
    # 增加 batch 维度 → (1, C, H, W)
    image = np.expand_dims(image, axis=0)

    return image


if __name__ == "__main__":
    image_url = "onnx_export/zidane.jpg"

    image = cv2.imread(image_url)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    processed_image = preprocess_image(image)
    print(processed_image.shape)

simplify_tokenizer.py

import json
import regex as re
from typing import Dict, List, Tuple, Optional


def bytes_to_unicode() -> Dict[int, str]:
    """
    与 CLIP / GPT 系列相同的 bytes -> unicode 映射。
    """
    bs = (
        list(range(ord("!"), ord("~") + 1))
        + list(range(ord("¡"), ord("¬") + 1))
        + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))


def get_pairs(word: Tuple[str, ...]):
    """
    返回一个 word 中相邻 symbol 的 pair 集合。
    word 是一个由字符串组成的 tuple。
    """
    pairs = set()
    prev_char = word[0]
    for char in word[1:]:
        pairs.add((prev_char, char))
        prev_char = char
    return pairs


class SimpleCLIPBPETokenizer:
    """
    纯 Python 的 CLIP BPE tokenizer 简化实现：
      - 从 vocab.json / merges.txt 加载词表与 BPE 规则
      - 实现 bytes_to_unicode + BPE + 正则分词
      - 提供 encode(text) -> (input_ids, attention_mask)

    用法示例：
        tokenizer = SimpleCLIPBPETokenizer(
            vocab_file="models/vocab.json",
            merges_file="models/merges.txt",
            max_length=32,
            bos_token_id=49406,
            eos_token_id=49407,
        )
        ids, mask = tokenizer.encode("ear")
    """

    def __init__(
        self,
        vocab_file: str,
        merges_file: str,
        max_length: int = 32,
        bos_token_id: int = 49406,
        eos_token_id: int = 49407,
        added_tokens: Optional[Dict[str, int]] = None,
        bpe_vocab_size: int = 49152,  # 一般和 merges 文件里取的行数一致
        do_lower_case: bool = True,
    ) -> None:
        self.max_length = max_length
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.do_lower_case = do_lower_case

        # 特殊 token 映射（可以通过入参覆盖）
        if added_tokens is None:
            added_tokens = {
                "<|startoftext|>": bos_token_id,
                "<|endoftext|>": eos_token_id,
            }
        self.added_tokens: Dict[str, int] = added_tokens
        self.unk_token: str = "<|endoftext|>"  # CLIP 里就是用这个当 unk

        # 正则分词规则
        self.pat = re.compile(
            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|"""
            r"""[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
            re.IGNORECASE,
        )

        # bytes -> unicode 映射
        self.byte_encoder = bytes_to_unicode()

        # 读取 merges，构建 bpe_ranks
        with open(merges_file, encoding="utf-8") as f:
            # 跳过第一行 "## merges" 之类，从第 2 行开始
            merges = f.read().strip().split("\n")[1 : bpe_vocab_size - 256 - 2 + 1]
        merges_pairs = [tuple(m.split()) for m in merges]
        self.bpe_ranks: Dict[Tuple[str, str], int] = {
            pair: i for i, pair in enumerate(merges_pairs)
        }

        # 读取 vocab.json
        with open(vocab_file, encoding="utf-8") as f:
            self.encoder: Dict[str, int] = json.load(f)

    # ----------------- 核心 BPE -----------------

    def _bpe(self, token: str) -> str:
        """
        单个 token 的 BPE 编码，返回空格分隔的子词字符串。
        """
        bpe_ranks = self.bpe_ranks

        word = tuple(token[:-1]) + (token[-1] + "</w>",)
        pairs = get_pairs(word)
        if not pairs:
            return token + "</w>"

        while True:
            bigram = min(pairs, key=lambda pair: bpe_ranks.get(pair, float("inf")))
            if bigram not in bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                except ValueError:
                    new_word.extend(word[i:])
                    break
                else:
                    new_word.extend(word[i:j])
                    i = j

                if (
                    word[i] == first
                    and i < len(word) - 1
                    and word[i + 1] == second
                ):
                    new_word.append(first + second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1

            word = tuple(new_word)
            if len(word) == 1:
                break
            pairs = get_pairs(word)

        return " ".join(word)

    # ----------------- 分词与编码 -----------------

    def _tokenize(self, text: str) -> List[str]:
        """
        文本 -> BPE token（字符串）列表。
        """
        if self.do_lower_case:
            text = text.lower()

        bpe_tokens: List[str] = []

        for token in self.pat.findall(text):
            # bytes -> unicode string
            token_encoded = "".join(
                self.byte_encoder[b] for b in token.encode("utf-8")
            )
            bpe_out = self._bpe(token_encoded)
            bpe_tokens.extend(bpe_out.split(" "))

        return bpe_tokens

    def _token_to_id(self, token: str) -> int:
        if token in self.added_tokens:
            return self.added_tokens[token]
        # unk 用 eos 替代（与原实现对齐）
        return self.encoder.get(token, self.encoder.get(self.unk_token))

    # ----------------- 对外接口：encode -----------------

    def encode(self, text: str) -> Tuple[List[int], List[int]]:
        """
        将文本编码为：
          - input_ids: 长度 max_length 的 id 序列
          - attention_mask: 同长度 mask（有效位置为 1，padding 为 0）
        """
        # 1. 先 BPE 分词
        bpe_tokens = self._tokenize(text)

        # 2. token -> id（此时不含 BOS/EOS）
        ids = [self._token_to_id(tok) for tok in bpe_tokens]
        len_ids = len(ids)

        # 3. 加 BOS / EOS
        num_special = 2  # BOS + EOS
        total_len = len_ids + num_special

        ids = [self.bos_token_id] + ids + [self.eos_token_id]

        # 4. 截断到 max_length
        ids = ids[: self.max_length]

        # 5. 有效长度（包含 BOS/EOS），与 HF attention_mask 逻辑对齐
        valid_len = min(total_len, self.max_length)

        # 6. 不足补 pad（用 eos_token_id 填充，CLIP 就是这么做的）
        if len(ids) < self.max_length:
            ids = ids + [self.eos_token_id] * (self.max_length - len(ids))

        # 7. attention_mask：前 valid_len 为 1，后面为 0
        attention_mask = [1] * valid_len + [0] * (self.max_length - valid_len)

        return ids, attention_mask


if __name__ == "__main__":
    # 简单测试
    vocab_file = "models/vocab.json"
    merges_file = "models/merges.txt"

    tokenizer = SimpleCLIPBPETokenizer(
        vocab_file=vocab_file,
        merges_file=merges_file,
        max_length=32,
        bos_token_id=49406,
        eos_token_id=49407,
        bpe_vocab_size=49152,
    )

    prompt = "person"
    ids, mask = tokenizer.encode(prompt)
    print("ids:", ids)
    print("len(ids):", len(ids))
    print("mask:", mask)
    print("sum(mask):", sum(mask))

detect_postprocess.py

import cv2
import numpy as np
import random

def process_sam3_results(
    outputs,
    img_h,
    img_w,
    score_thr=0.4,
    mask_thr=0.5,
    max_inst=30,
    boxes_normalized=True,
):
    """
    处理 SAM3 输出，返回结构化的检测结果列表。
    
    outputs: [pred_masks, pred_boxes, pred_logits]
    img_h, img_w: 原图高宽，用于缩放 mask 和 box
    """
    pred_masks, pred_boxes, pred_logits = outputs

    # 去掉 batch 维
    if pred_masks.ndim == 4:
        pred_masks = pred_masks[0]      # [N, Hm, Wm]
    if pred_boxes.ndim == 3:
        pred_boxes = pred_boxes[0]      # [N, 4]
    if pred_logits.ndim == 2:
        pred_logits = pred_logits[0]    # [N]

    pred_masks  = pred_masks.astype(np.float32)
    pred_boxes  = pred_boxes.astype(np.float32)
    pred_logits = pred_logits.astype(np.float32)

    N, Hm, Wm = pred_masks.shape

    # logits -> scores
    scores = 1.0 / (1.0 + np.exp(-pred_logits))  # [N]

    # 排序 + 阈值 + topk
    indices = list(range(N))
    indices = sorted(indices, key=lambda i: float(scores[i]), reverse=True)
    indices = [i for i in indices if float(scores[i]) >= score_thr]
    indices = indices[:max_inst]

    results = []
    for i in indices:
        mask  = pred_masks[i]          # [Hm, Wm]
        score = float(scores[i])
        box   = pred_boxes[i]          # [4]

        # 1) 先把 mask resize 到原图大小
        mask_resized = cv2.resize(mask, (img_w, img_h), interpolation=cv2.INTER_LINEAR)
        m = (mask_resized > mask_thr).astype(np.uint8)  # [H, W]

        if m.sum() == 0:
            continue

        # 2) 处理框坐标
        x1, y1, x2, y2 = box

        if boxes_normalized:
            x1 = int(x1 * img_w)
            x2 = int(x2 * img_w)
            y1 = int(y1 * img_h)
            y2 = int(y2 * img_h)
        else:
            x1 = int(x1)
            y1 = int(y1)
            x2 = int(x2)
            y2 = int(y2)

        # clamp 一下防止越界
        x1 = max(0, min(x1, img_w - 1))
        x2 = max(0, min(x2, img_w - 1))
        y1 = max(0, min(y1, img_h - 1))
        y2 = max(0, min(y2, img_h - 1))

        results.append({
            "mask": m,          # [H, W] uint8
            "box": [x1, y1, x2, y2],
            "score": score,
        })
        
    return results

def draw_sam3_results(
    img,
    results,
):
    """
    在原图上画检测结果。
    
    img: 原图 (H, W, 3)
    results: process_sam3_results 返回的列表
    """
    vis_img = img.copy().astype(np.float32)
    img_h, img_w = img.shape[:2]

    for item in results:
        m = item["mask"]
        box = item["box"]
        score = item["score"]

        # 随机颜色
        color = [random.randint(0, 255) for _ in range(3)]

        # 3 通道 mask
        m3 = np.stack([m, m, m], axis=-1)  # [H, W, 3]

        # 在原图上叠加半透明颜色
        vis_img = np.where(
            m3 == 1,
            vis_img * 0.5 + np.array(color, dtype=np.float32) * 0.5,
            vis_img,
        )

        x1, y1, x2, y2 = box

        # 画框
        cv2.rectangle(
            vis_img,
            (x1, y1),
            (x2, y2),
            color,
            2,
            lineType=cv2.LINE_AA,
        )

        # 写分数
        cv2.putText(
            vis_img,
            f"{score:.2f}",
            (x1, max(y1 - 5, 0)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            color,
            1,
            cv2.LINE_AA,
        )

    vis_img = np.clip(vis_img, 0, 255).astype(np.uint8)
    return vis_img

main.py

import numpy as np
import cv2
import onnxruntime as ort
from image_preprocess import preprocess_image
from simplify_tokenizer import SimpleCLIPBPETokenizer
from detect_postprocess import process_sam3_results, draw_sam3_results


if __name__ == "__main__":
    image_url = "dog.jpg"
    onnx_file_path = "sam3.onnx"

    image = cv2.imread(image_url)

    session = ort.InferenceSession(
        onnx_file_path,
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
    )

    vocab_file = "vocab.json"
    merges_file = "merges.txt"
    prompt = "dog"

    tokenizer = SimpleCLIPBPETokenizer(
        vocab_file=vocab_file,
        merges_file=merges_file,
        max_length=32,
        bos_token_id=49406,
        eos_token_id=49407,
        bpe_vocab_size=49152,
    )

    processed_image = preprocess_image(image)
    
    ids, mask = tokenizer.encode(prompt)
    input_ids = np.array(ids, dtype=np.int64).reshape(1, -1)
    attention_mask = np.array(mask, dtype=np.int64).reshape(1, -1)

    input_dict = {
        "pixel_values": processed_image.astype("float32"),  #(1, 3, 1008, 1008)
        "input_ids": input_ids,                             #(1, 32)     
        "attention_mask": attention_mask,                   #(1, 32)    
    }

    outputs = session.run(None, input_dict)

    image = np.array(image)
    results = process_sam3_results(
        outputs,
        img_h=image.shape[0],
        img_w=image.shape[1],
        score_thr=0.6,
        mask_thr=0.5,
        max_inst=30,
        boxes_normalized=True,
    )
    vis_img = draw_sam3_results(image, results)
    
    cv2.imwrite("result.jpg", vis_img)

tensorrt推理

common.py

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
import os
import ctypes
from typing import Optional, List

import numpy as np
import tensorrt as trt
from cuda import cuda, cudart

try:
    # Sometimes python does not understand FileNotFoundError
    FileNotFoundError
except NameError:
    FileNotFoundError = IOError

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def check_cuda_err(err):
    if isinstance(err, cuda.CUresult):
        if err != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError("Cuda Error: {}".format(err))
    if isinstance(err, cudart.cudaError_t):
        if err != cudart.cudaError_t.cudaSuccess:
            raise RuntimeError("Cuda Runtime Error: {}".format(err))
    else:
        raise RuntimeError("Unknown error type: {}".format(err))

def cuda_call(call):
    err, res = call[0], call[1:]
    check_cuda_err(err)
    if len(res) == 1:
        res = res[0]
    return res

def GiB(val):
    return val * 1 << 30


def add_help(description):
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    args, _ = parser.parse_known_args()


def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""):
    """
    Parses sample arguments.

    Args:
        description (str): Description of the sample.
        subfolder (str): The subfolder containing data relevant to this sample
        find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.

    Returns:
        str: Path of data directory.
    """

    # Standard command-line arguments for all samples.
    kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "-d",
        "--datadir",
        help="Location of the TensorRT sample data directory, and any additional data directories.",
        action="append",
        default=[kDEFAULT_DATA_ROOT],
    )
    args, _ = parser.parse_known_args()

    def get_data_path(data_dir):
        # If the subfolder exists, append it to the path, otherwise use the provided path as-is.
        data_path = os.path.join(data_dir, subfolder)
        if not os.path.exists(data_path):
            if data_dir != kDEFAULT_DATA_ROOT:
                print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")
            data_path = data_dir
        # Make sure data directory exists.
        if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:
            print(
                "WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(
                    data_path
                )
            )
        return data_path

    data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
    return data_paths, locate_files(data_paths, find_files, err_msg)


def locate_files(data_paths, filenames, err_msg=""):
    """
    Locates the specified files in the specified data directories.
    If a file exists in multiple data directories, the first directory is used.

    Args:
        data_paths (List[str]): The data directories.
        filename (List[str]): The names of the files to find.

    Returns:
        List[str]: The absolute paths of the files.

    Raises:
        FileNotFoundError if a file could not be located.
    """
    found_files = [None] * len(filenames)
    for data_path in data_paths:
        # Find all requested files.
        for index, (found, filename) in enumerate(zip(found_files, filenames)):
            if not found:
                file_path = os.path.abspath(os.path.join(data_path, filename))
                if os.path.exists(file_path):
                    found_files[index] = file_path

    # Check that all files were found
    for f, filename in zip(found_files, filenames):
        if not f or not os.path.exists(f):
            raise FileNotFoundError(
                "Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg)
            )
    return found_files


class HostDeviceMem:
    """Pair of host and device memory, where the host memory is wrapped in a numpy array"""
    def __init__(self, size: int, dtype: np.dtype):
        nbytes = size * dtype.itemsize
        host_mem = cuda_call(cudart.cudaMallocHost(nbytes))
        pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))

        self._host = np.ctypeslib.as_array(ctypes.cast(host_mem, pointer_type), (size,))
        self._device = cuda_call(cudart.cudaMalloc(nbytes))
        self._nbytes = nbytes

    @property
    def host(self) -> np.ndarray:
        return self._host

    @host.setter
    def host(self, arr: np.ndarray):
        if arr.size > self.host.size:
            raise ValueError(
                f"Tried to fit an array of size {arr.size} into host memory of size {self.host.size}"
            )
        np.copyto(self.host[:arr.size], arr.flat, casting='safe')

    @property
    def device(self) -> int:
        return self._device

    @property
    def nbytes(self) -> int:
        return self._nbytes

    def __str__(self):
        return f"Host:\n{self.host}\nDevice:\n{self.device}\nSize:\n{self.nbytes}\n"

    def __repr__(self):
        return self.__str__()

    def free(self):
        cuda_call(cudart.cudaFree(self.device))
        cuda_call(cudart.cudaFreeHost(self.host.ctypes.data))


# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
# If engine uses dynamic shapes, specify a profile to find the maximum input & output size.
def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda_call(cudart.cudaStreamCreate())
    tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
    for binding in tensor_names:
        # get_tensor_profile_shape returns (min_shape, optimal_shape, max_shape)
        # Pick out the max shape to allocate enough memory for the binding.
        shape = engine.get_tensor_shape(binding) if profile_idx is None else engine.get_tensor_profile_shape(binding, profile_idx)[-1]
        shape_valid = np.all([s >= 0 for s in shape])
        if not shape_valid and profile_idx is None:
            raise ValueError(f"Binding {binding} has dynamic shape, " +\
                "but no profile was specified.")
        size = trt.volume(shape)
        if engine.has_implicit_batch_dimension:
            size *= engine.max_batch_size
        dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(binding)))

        # Allocate host and device buffers
        bindingMemory = HostDeviceMem(size, dtype)

        # Append the device buffer to device bindings.
        bindings.append(int(bindingMemory.device))

        # Append to the appropriate list.
        if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
            inputs.append(bindingMemory)
        else:
            outputs.append(bindingMemory)
    return inputs, outputs, bindings, stream


# Frees the resources allocated in allocate_buffers
def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem], stream: cudart.cudaStream_t):
    for mem in inputs + outputs:
        mem.free()
    cuda_call(cudart.cudaStreamDestroy(stream))


# Wrapper for cudaMemcpy which infers copy size and does error checking
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
    nbytes = host_arr.size * host_arr.itemsize
    cuda_call(cudart.cudaMemcpy(device_ptr, host_arr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))


# Wrapper for cudaMemcpy which infers copy size and does error checking
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
    nbytes = host_arr.size * host_arr.itemsize
    cuda_call(cudart.cudaMemcpy(host_arr, device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))


def _do_inference_base(inputs, outputs, stream, execute_async):
    # Transfer input data to the GPU.
    kind = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
    [cuda_call(cudart.cudaMemcpyAsync(inp.device, inp.host, inp.nbytes, kind, stream)) for inp in inputs]
    # Run inference.
    execute_async()
    # Transfer predictions back from the GPU.
    kind = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
    [cuda_call(cudart.cudaMemcpyAsync(out.host, out.device, out.nbytes, kind, stream)) for out in outputs]
    # Synchronize the stream
    cuda_call(cudart.cudaStreamSynchronize(stream))
    # Return only the host outputs.
    return [out.host for out in outputs]


def do_inference(context, engine, bindings, inputs, outputs, stream):
    def execute_async_func():
        context.execute_async_v3(stream_handle=stream)
    # Setup context tensor address.
    num_io = engine.num_io_tensors
    for i in range(num_io):
        context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
    return _do_inference_base(inputs, outputs, stream, execute_async_func)

main.py

import numpy as np
import cv2
import tensorrt as trt
import common
from image_preprocess import preprocess_image
from simplify_tokenizer import SimpleCLIPBPETokenizer
from detect_postprocess import process_sam3_results, draw_sam3_results


if __name__ == "__main__":
    image_url = "dog.jpg"
    engine_file_path = "sam3.engine"

    image = cv2.imread(image_url)

    logger = trt.Logger(trt.Logger.WARNING)
    with open(engine_file_path, "rb") as f, trt.Runtime(logger) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()
    inputs, outputs, bindings, stream = common.allocate_buffers(engine)

    vocab_file = "vocab.json"
    merges_file = "merges.txt"
    prompt = "dog"

    tokenizer = SimpleCLIPBPETokenizer(
        vocab_file=vocab_file,
        merges_file=merges_file,
        max_length=32,
        bos_token_id=49406,
        eos_token_id=49407,
        bpe_vocab_size=49152,
    )

    processed_image = preprocess_image(image)
    
    ids, mask = tokenizer.encode(prompt)
    input_ids = np.array(ids, dtype=np.int64).reshape(1, -1)
    attention_mask = np.array(mask, dtype=np.int64).reshape(1, -1)

    np.copyto(inputs[0].host, processed_image.ravel())
    np.copyto(inputs[1].host, input_ids.ravel())
    np.copyto(inputs[2].host, attention_mask.ravel())

    output = common.do_inference(context, engine=engine, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
    outputs = [output[0].reshape(1, 200, 288, 288), output[1].reshape(1, 200, 4), output[2].reshape(1, 200)]

    image = np.array(image)
    results = process_sam3_results(
        outputs,
        img_h=image.shape[0],
        img_w=image.shape[1],
        score_thr=0.6,
        mask_thr=0.5,
        max_inst=30,
        boxes_normalized=True,
    )
    vis_img = draw_sam3_results(image, results)
    
    cv2.imwrite("result.jpg", vis_img)

结果如下

代码和图片见：https://github.com/taifyang/sam-inference