chore: import upstream snapshot with attribution

samples/python/dds_faster_rcnn/README.md

@@ -0,0 +1,105 @@
+# DDS Faster R-CNN Object Detection in TensorRT
+## Introduction
+The `dds_faster_rcnn` sample demonstrates the usage of [tensorrt.IOutputAllocator](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/ExecutionContext.html#tensorrt.IOutputAllocator) in TensorRT to execute networks with data-dependent shape (DDS) outputs. In this sample, we showcase an end-to-end workflow for building and running an object detection model [Faster-RCNN](https://arxiv.org/abs/1506.01497).
+
+### What are Data-Dependent Shapes (DDS)?
+Data-Dependent Shapes (DDS) refer to shapes of layer outputs in a neural network which depend on the input data to the layer; in other words, it cannot be inferred solely by inspecting the shapes of the layer's input tensors.  An example of this is the output shape of the `INonZeroLayer`, which is determined by the number of non-zero elements in the input tensor.
+
+DDS outputs are common in models that involve dynamic processing, such as object detection, segmentation, and natural language processing.
+
+### What is an `IOutputAllocator`?
+An `IOutputAllocator` is an interface in TensorRT that defines a class responsible for dynamically allocating and managing the device memory for output tensors of a TensorRT engine. The class implementing this interface must provide a way to allocate and deallocate memory for output tensors, which can vary in size depending on the input data.
+
+### Why do we need to implement `IOutputAllocator`
+In traditional models, the output shapes are typically fixed and known at build time. However, in the case of data-dependent shaped (DDS) outputs, the output size is only known at inference time. This means that the memory allocation for output tensors cannot be determined until the model is actually run with a specific input. To handle this situation, TensorRT provides the `IOutputAllocator` interface, which allows developers to implement a custom memory allocation strategy for DDS outputs. By implementing this interface, developers can ensure that the output tensors are properly allocated and deallocated during inference, avoiding potential memory issues and improving the overall performance of the model.
+
+### How does `IOutputAllocator` work?
+To implement the `IOutputAllocator` interface, you need to provide implementations for the following two key methods:
+
+- `reallocate_output_async(self, tensor_name, memory, size, alignment, stream)`: This method is responsible for allocating or reallocating memory for an output tensor. It is called during the inference phase when the output tensor size is known. The method takes in parameters such as the tensor name, current memory address, new size, alignment, and CUDA stream, and returns the new memory address.
+- `notify_shape(self, tensor_name, shape)`: This method is used to notify the allocator of a change in the shape of an output tensor. It is typically called after reallocate_output_async() to update the allocator's internal state with the new shape information.
+During inference, the TensorRT engine will call these methods to manage the memory allocation for DDS output tensors. The `IOutputAllocator` implementation is responsible for ensuring that the memory allocation is properly handled, taking into account factors such as memory fragmentation, alignment, and performance optimization.
+
+Here is a high-level overview of the workflow:
+
+1. Instantiate the output allocator and attach to TensorRT with [IExecutionContext.set_output_allocator()](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/ExecutionContext.html#tensorrt.IExecutionContext.set_output_allocator)
+1. The TensorRT engine determines that an output tensor needs to be allocated or reallocated.
+1. `reallocate_output_async` is called to allocate or reallocate memory for the output tensor.
+1. The allocator updates its internal state and returns the new memory address.
+1. The TensorRT engine uses the new memory address to store the output tensor data.
+1. `notify_shape()` method is called to update the allocator's internal state with the new shape information.
+
+By implementing the `IOutputAllocator` interface, developers can create custom memory allocation strategies that optimize performance, reduce memory fragmentation, and improve the overall efficiency of their model inference.
+
+## Setup
+We recommend running these scripts on an environment with TensorRT >= 10.8.0.
+
+Install TensorRT as per the [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html). You will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the `python3-libnvinfer` and `python3-libnvinfer-dev` packages on your TensorRT download.
+
+To simplify TensorRT installation, use an NGC Docker Image, such as:
+
+```bash
+docker pull nvcr.io/nvidia/tensorrt:25.01-py3
+```
+
+Install all dependencies listed in requirements.txt:
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## Model Conversion
+To start, download the pre-trained Faster R-CNN model in ONNX format using the following command:
+
+```bash
+wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-12.onnx
+```
+
+With the ONNX model downloaded, run the following command to prepare it for TensorRT engine conversion:
+
+```bash
+python3 modify_onnx.py \
+    --input ./FasterRCNN-12.onnx \
+    --output ./fasterrcnn12_trt.onnx
+```
+
+This will create a modified ONNX graph file that is ready for conversion to a TensorRT engine.
+
+## Build TensorRT Engine
+
+To build the TensorRT engine, run the following command:
+
+```bash
+python3 build_engine.py \
+    --onnx ./fasterrcnn12_trt.onnx \
+    --engine ./fasterrcnn12_trt.engine
+```
+
+## Inference
+To test the built TensorRT engine, download a test image using the following command:
+
+```bash
+wget https://onnxruntime.ai/images/demo.jpg
+```
+
+Then, run the inference script using the following command:
+
+```
+python3 infer.py \
+    --engine ./fasterrcnn12_trt.engine \
+    --input ./demo.jpg \
+    --output ./output_dir \
+    --labels labels_coco_80.txt
+```
+This will perform object detection on the test image and save the output to the specified directory (`output_dir` in this case).
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+February 2025
+Initial release

samples/python/dds_faster_rcnn/build_engine.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 os
+import sys
+import logging
+import argparse
+
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("EngineBuilder").setLevel(logging.INFO)
+log = logging.getLogger("EngineBuilder")
+
+
+class EngineBuilder:
+    """
+    Parses an ONNX graph and builds a TensorRT engine from it.
+    """
+
+    def __init__(self, verbose=False, workspace=8):
+        """
+        :param verbose: If enabled, a higher verbosity level will be set on the TensorRT logger.
+        :param workspace: Max memory workspace to allow, in Gb.
+        """
+        self.trt_logger = trt.Logger(trt.Logger.INFO)
+        if verbose:
+            self.trt_logger.min_severity = trt.Logger.Severity.VERBOSE
+
+        trt.init_libnvinfer_plugins(self.trt_logger, namespace="")
+
+        self.builder = trt.Builder(self.trt_logger)
+        self.config = self.builder.create_builder_config()
+        one_GiB = 2**30
+        self.config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace * one_GiB)
+        self.network = None
+        self.parser = None
+
+    def create_network(self, onnx_path):
+        """
+        Parse the ONNX graph and create the corresponding TensorRT network definition.
+        :param onnx_path: The path to the ONNX graph to load.
+        """
+
+        self.network = self.builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        self.parser = trt.OnnxParser(self.network, self.trt_logger)
+
+        onnx_path = os.path.realpath(onnx_path)
+        with open(onnx_path, "rb") as f:
+            if not self.parser.parse(f.read()):
+                for error in range(self.parser.num_errors):
+                    log.error(self.parser.get_error(error))
+                raise RuntimeError(
+                    f"Failed to load ONNX file: {onnx_path}. Check the logs for more details or run with --verbose."
+                )
+
+        log.info("Network Description")
+
+        profile = self.builder.create_optimization_profile()
+        profile.set_shape("image", min=(3, 1, 1), opt=(3, 800, 800), max=(3, 800, 1312))
+        self.config.add_optimization_profile(profile)
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+        for input in inputs:
+            log.info(f"Input '{input.name}' with shape {input.shape} and dtype {input.dtype}")
+        outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)]
+        for output in outputs:
+            log.info(f"Output '{output.name}' with shape {output.shape} and dtype {output.dtype}")
+
+    def create_engine(
+        self,
+        engine_path,
+    ):
+        """
+        Build the TensorRT engine and serialize it to disk.
+        :param engine_path: The path where to serialize the engine to.
+        """
+        engine_path = os.path.realpath(engine_path)
+        engine_dir = os.path.dirname(engine_path)
+        os.makedirs(engine_dir, exist_ok=True)
+        log.info(f"Building Engine in {engine_path}")
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+
+        log.info(f"Reading timing cache from file: {args.timing_cache}")
+        common.setup_timing_cache(self.config, args.timing_cache)
+
+        engine_bytes = self.builder.build_serialized_network(self.network, self.config)
+        if engine_bytes is None:
+            raise RuntimeError("Failed to create engine. Check the logs for more details or run with --verbose.")
+
+        log.info(f"Serializing timing cache to file: {args.timing_cache}")
+        common.save_timing_cache(self.config, args.timing_cache)
+
+        with open(engine_path, "wb") as f:
+            log.info(f"Serializing engine to file: {engine_path}")
+            f.write(engine_bytes)
+
+
+def main(args):
+    builder = EngineBuilder(args.verbose, args.workspace)
+    builder.create_network(args.onnx)
+    builder.create_engine(
+        args.engine,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--onnx", required=True, help="The input ONNX model file to load")
+    parser.add_argument("-e", "--engine", required=True, help="The output path for the TRT engine")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable more verbose log output")
+    parser.add_argument(
+        "-w",
+        "--workspace",
+        default=8,
+        type=int,
+        help="The max memory workspace size to allow in Gb, default: 8",
+    )
+    parser.add_argument(
+        "--timing_cache",
+        default="./timing.cache",
+        help="The file path for timing cache, default: ./timing.cache",
+    )
+    args = parser.parse_args()
+    main(args)

samples/python/dds_faster_rcnn/infer.py

@@ -0,0 +1,485 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 os
+import sys
+import time
+import argparse
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import driver as cuda, runtime as cudart
+from PIL import Image
+from pathlib import Path
+import threading
+from visualize import visualize_detections
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+from common import cuda_call
+
+
+class AllocatorState:
+    """
+    Represents the state of an allocator for a tensor.
+    """
+
+    def __init__(self, ptr, size, dim=None):
+        """
+        :param ptr: The pointer to the allocated device memory.
+        :param size: The size of the allocated device memory.
+        :param dim: The dimensions of the tensor.
+        """
+        self.ptr = ptr
+        self.size = size
+        self.dim = dim
+        self.lock = threading.Lock()
+
+    def update(self, ptr=None, size=None, dims=None):
+        """
+        Updates the state of the allocator.
+
+        :param ptr: The new pointer to the allocated device memory. If None, the current pointer is not changed.
+        :param size: The new size of the allocated device memory. If None, the current size is not changed.
+        :param dims: The new dimensions of the tensor. If None, the current dimensions are not changed.
+        """
+        with self.lock:
+            if ptr is not None:
+                self.ptr = ptr
+            if size is not None:
+                self.size = size
+            if dims is not None:
+                self.dims = dims
+
+
+class MyOutputAllocator(trt.IOutputAllocator):
+    """
+    Custom output allocator class.
+    """
+
+    def __init__(self, verbose=False):
+        """
+        :param verbose: If True, enables verbose logging.
+        """
+        trt.IOutputAllocator.__init__(self)
+
+        self.lock = threading.Lock()
+        self.states = {}
+        self.verbose = verbose
+
+    def reallocate_output_async(self, tensor_name, current_memory, size, alignment, stream):
+        """
+        Reallocates output memory for the given tensor.
+
+        :param tensor_name: The name of the tensor.
+        :param current_memory: The current device memory pointer.
+        :param size: The new size of the device memory block.
+        :param alignment: The alignment of the device memory block.
+        :param stream: The CUDA stream.
+        :return: The new memory pointer.
+        """
+        size = max(size, 1)
+        ptr = current_memory
+        with self.lock:
+            if tensor_name not in self.states or size > self.states[tensor_name].size:
+                ptr = cuda_call(cudart.cudaMalloc(size))
+                if tensor_name in self.states:
+                    cuda_call(cudart.cudaFree(self.states[tensor_name].ptr))
+                    self.states[tensor_name].update(ptr=ptr, size=size)
+                else:
+                    self.states[tensor_name] = AllocatorState(ptr=ptr, size=size)
+                if self.verbose:
+                    print(f"Reallocated {size} bytes for tensor '{tensor_name}' to {ptr}")
+        return ptr
+
+    def notify_shape(self, tensor_name, dims):
+        """
+        Notifies the allocator of a change in the shape of the tensor.
+
+        :param tensor_name: The name of the tensor.
+        :param dims: The new dimensions of the tensor.
+        """
+        with self.lock:
+            assert tensor_name in self.states, f'Tensor "{tensor_name}" is not in states.'
+            self.states[tensor_name].update(dims=dims)
+            if self.verbose:
+                print(f"Updated shape for tensor '{tensor_name}': {dims}")
+
+    def __del__(self):
+        try:
+            with self.lock:
+                for tensor_name, item in self.states.items():
+                    if item.ptr is not None:
+                        cuda_call(cudart.cudaFree(item.ptr))
+                        if self.verbose:
+                            print(f"Freed memory for tensor '{tensor_name}'")
+                self.states.clear()
+        except Exception:
+            # Silently handle cleanup failures to prevent exceptions during object deletion
+            pass
+
+
+class PoolAllocator(trt.IGpuAsyncAllocator):
+    """
+    A custom GPU async allocator class that manages memory allocation and deallocation.
+
+    It utilizes the CUDA memory pool API to optimize memory allocation and minimize fragmentation.
+    """
+
+    def __init__(self):
+        """
+        Initializes the PoolAllocator instance.
+
+        Creates a CUDA memory pool with the specified properties and sets the release threshold to the maximum possible value.
+        """
+        trt.IGpuAsyncAllocator.__init__(self)
+
+        pool_props = cudart.cudaMemPoolProps()
+        pool_props.allocType = cudart.cudaMemAllocationType.cudaMemAllocationTypePinned
+        pool_props.handleTypes = cudart.cudaMemAllocationHandleType.cudaMemHandleTypeNone
+        pool_props.location.type = cudart.cudaMemLocationType.cudaMemLocationTypeDevice
+        pool_props.location.id = 0
+
+        self.pool = cuda_call(cudart.cudaMemPoolCreate(pool_props))
+
+        max_threshold = np.uint64(np.iinfo(np.uint64).max)
+        cuda_call(
+            cudart.cudaMemPoolSetAttribute(
+                self.pool, cudart.cudaMemPoolAttr.cudaMemPoolAttrReleaseThreshold, cuda.cuuint64_t(max_threshold)
+            )
+        )
+
+    def allocate_async(self, size: int, alignment: int, flags: int, stream: cudart.cudaStream_t):
+        """
+        Allocates memory asynchronously from the CUDA memory pool.
+
+        :param size: The size of the memory block to allocate.
+        :param alignment: The alignment of the memory block.
+        :param flags: The flags for the allocation.
+        :param stream: The CUDA stream for the allocation.
+        :return: The pointer to the allocated device memory.
+        """
+        ptr = cuda_call(cudart.cudaMallocFromPoolAsync(size, self.pool, stream))
+        return ptr
+
+    def deallocate_async(self, memory, stream: cudart.cudaStream_t):
+        """
+        Deallocates memory asynchronously.
+
+        :param memory: The pointer to the memory to deallocate.
+        :param stream: The CUDA stream for the deallocation.
+        :return: True if the deallocation was successful.
+        """
+        cuda_call(cudart.cudaFreeAsync(memory, stream))
+        return True
+
+    def __del__(self):
+        try:
+            if self.pool:
+                cuda_call(cudart.cudaMemPoolDestroy(self.pool))
+        except Exception:
+            # Silently handle cleanup failures to prevent exceptions during object deletion
+            pass
+
+
+class TensorRTInfer:
+    """
+    Implements inference for the FasterRCNN TensorRT engine.
+    """
+
+    def __init__(self, engine_path, use_custom_gpu_allocator=False, verbose=False):
+        """
+        Initializes the TensorRTInfer instance.
+
+        :param engine_path: The path to the serialized engine to load from disk.
+        :param use_custom_gpu_allocator: If True, uses a custom GPU allocator.
+        :param verbose: If True, enables verbose logging.
+        """
+        # Load TRT engine
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        if verbose:
+            self.logger.min_severity = trt.Logger.VERBOSE
+        trt.init_libnvinfer_plugins(self.logger, namespace="")
+        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
+            assert runtime
+            if use_custom_gpu_allocator:
+                self.my_pool_allocator = PoolAllocator()
+                runtime.gpu_allocator = self.my_pool_allocator
+            self.engine = runtime.deserialize_cuda_engine(f.read())
+        assert self.engine
+        self.context = self.engine.create_execution_context()
+        assert self.context
+
+        self.my_output_allocator = MyOutputAllocator(verbose=True)
+        # Setup I/O bindings
+        self.inputs = []
+        self.outputs = []
+        self.allocations = []
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            is_input = False
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                is_input = True
+            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
+
+            # trt.nptype returns a python 'type'. For here we want a numpy 'dtype' object
+            # instead to get more info about the dtype (dtype.itemsize in this case)
+            dtype = np.dtype(dtype)
+            shape = self.engine.get_tensor_shape(name)
+
+            # Use the max shape in the profile for dynamic shaped inputs
+            if is_input and any(value for value in shape if value < 0):
+                assert self.engine.num_optimization_profiles > 0
+                profile_shape = self.engine.get_tensor_profile_shape(name, 0)
+                assert len(profile_shape) == 3  # min,opt,max
+                # Set the *max* profile as binding shape
+                shape = profile_shape[2]
+
+            if is_input:
+                nbytes = np.prod(shape) * dtype.itemsize
+                allocation = cuda_call(cudart.cudaMalloc(nbytes))
+            else:
+                self.context.set_output_allocator(name, self.my_output_allocator)
+                allocation = cuda_call(
+                    cudart.cudaMalloc(128 * dtype.itemsize)
+                )  # Random number. More will be allocated using our custom allocator
+
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": dtype,
+                "shape": list(shape),
+                "allocation": allocation,
+            }
+            self.allocations.append(allocation)
+            if is_input:
+                self.inputs.append(binding)
+            else:
+                self.outputs.append(binding)
+            print(
+                f"{'Input' if is_input else 'Output'} '{binding['name']}' with shape {binding['shape']} and dtype {binding['dtype']}"
+            )
+
+        assert len(self.inputs) > 0
+        assert len(self.outputs) > 0
+        assert len(self.allocations) > 0
+
+    def input_spec(self):
+        """
+        Get the specs for the input tensor of the network. Useful to prepare memory allocations.
+        :return: Two items, the shape of the input tensor and its (numpy) datatype.
+        """
+        return self.inputs[0]["shape"], self.inputs[0]["dtype"]
+
+    def output_spec(self):
+        """
+        Get the specs for the output tensors of the network. Useful to prepare memory allocations.
+        :return: A list with two items per element, the shape and (numpy) datatype of each output tensor.
+        """
+        specs = []
+        for o in self.outputs:
+            specs.append((o["shape"], o["dtype"]))
+        return specs
+
+    def preprocess_image(self, image):
+        """
+        Preprocesses an image for inference. See also
+        https://github.com/onnx/models/tree/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn#preprocessing-steps
+
+        :param image: The image to preprocess.
+        :return: The preprocessed image as a numpy array.
+        """
+        ratio = 800.0 / min(image.size[0], image.size[1])
+        image = image.resize((int(ratio * image.size[0]), int(ratio * image.size[1])), Image.BILINEAR)
+
+        # RGB -> BGR
+        image = np.array(image)[:, :, [2, 1, 0]].astype("float32")
+
+        # HWC -> CHW
+        image = np.transpose(image, [2, 0, 1])
+
+        # Normalize
+        mean_vec = np.array([102.9801, 115.9465, 122.7717])
+        for i in range(image.shape[0]):
+            image[i, :, :] = image[i, :, :] - mean_vec[i]
+
+        # Pad to be divisible of 32
+        padded_h = int(np.ceil(image.shape[1] / 32) * 32)
+        padded_w = int(np.ceil(image.shape[2] / 32) * 32)
+
+        padded_image = np.zeros((3, padded_h, padded_w), dtype=np.float32)
+        padded_image[:, : image.shape[1], : image.shape[2]] = image
+        image = padded_image
+
+        return image
+
+    def infer(self, arr):
+        """
+        Execute inference on an image.
+
+        :param arr: A numpy array for the input image values.
+        :return A list of outputs as numpy arrays.
+        """
+        # Copy I/O and Execute
+        common.memcpy_host_to_device(self.inputs[0]["allocation"], arr)
+        self.context.execute_v2(self.allocations)
+
+        # copy outputs to host
+        return_outputs = []
+        for output in self.outputs:
+            final_shape = self.my_output_allocator.states[output["name"]].dims
+            host_arr = np.random.random(final_shape).astype(output["dtype"])
+            device_ptr = self.my_output_allocator.states[output["name"]].ptr
+
+            nbytes = np.prod(final_shape) * output["dtype"].itemsize
+            common.memcpy_device_to_host(host_arr, device_ptr)
+
+            return_outputs.append(host_arr)
+
+        return return_outputs
+
+    def process(self, arr):
+        """
+        Execute inference on an image. The image should already be preprocessed. Memory
+        copying to and from the GPU device will be performed here.
+
+        :param arr: A numpy array holding the image values.
+        :return: A list of detected object with box, score, class included.
+        """
+        preprocess_arr = self.preprocess_image(arr.copy())
+        self.context.set_input_shape("image", preprocess_arr.shape)
+
+        # Run inference
+        outputs = self.infer(preprocess_arr)
+
+        # Post-process the results
+        scale = 800.0 / min(arr.size[0], arr.size[1])
+
+        boxes = outputs[0]
+        labels = outputs[1]
+        scores = outputs[2]
+        num = len(labels)
+
+        detections = []
+        for i in range(num):
+            if scores[i] > 0.9:
+                detections.append(
+                    {
+                        "xmin": boxes[i][0] / scale,
+                        "ymin": boxes[i][1] / scale,
+                        "xmax": boxes[i][2] / scale,
+                        "ymax": boxes[i][3] / scale,
+                        "score": scores[i],
+                        "class": labels[i] - 1,
+                    }
+                )
+        return detections
+
+
+def main(args):
+    if args.output:
+        args.output.resolve().mkdir(exist_ok=True, parents=True)
+
+    labels = []
+    if args.labels:
+        with open(args.labels) as f:
+            for label in f:
+                labels.append(label.strip())
+
+    trt_infer = TensorRTInfer(args.engine, args.use_custom_gpu_allocator, args.verbose)
+    if args.input:
+        print(f"\nInferring data in {args.input}")
+        image_paths = []
+        if args.input.is_dir():
+            for p in args.input.iterdir():
+                image_paths.append(p)
+        else:
+            image_paths.append(args.input)
+
+        for image_path in image_paths:
+            image = Image.open(image_path)
+            detections = trt_infer.process(image)
+            if args.output:
+                # Image Visualizations
+                output_path = args.output / f"{image_path.stem}.png"
+                visualize_detections(image_path, output_path, detections, labels)
+
+                # Text Results
+                output_results = ""
+                for d in detections:
+                    line = [
+                        d["xmin"],
+                        d["ymin"],
+                        d["xmax"],
+                        d["ymax"],
+                        d["score"],
+                    ]
+                    output_results += "\t".join([str(f) for f in line]) + "\n"
+                with open(args.output / f"{image_path.stem}.txt", "w") as f:
+                    f.write(output_results)
+    else:
+        print("No input provided, running in benchmark mode")
+        shape, dtype = trt_infer.input_spec()
+        batch = 255 * np.random.rand(*shape).astype(dtype)
+        trt_infer.context.set_input_shape("image", (batch.shape))
+        iterations = 200
+        times = []
+        for i in range(20):  # GPU warmup iterations
+            trt_infer.infer(batch)
+        for i in range(iterations):
+            start = time.time()
+            trt_infer.infer(batch)
+            times.append(time.time() - start)
+            print(f"Iteration {i+1} / {iterations}", end="\r")
+        print("Benchmark results include time for H2D and D2H memory copies")
+        print(f"Average Latency: {1000 * np.average(times):.3f} ms")
+
+    print("\nFinished Processing")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-e",
+        "--engine",
+        default=None,
+        required=True,
+        help="The serialized TensorRT engine",
+    )
+    parser.add_argument("-i", "--input", default=None, type=Path, help="Path to the image or directory to process")
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        type=Path,
+        help="Directory where to save the visualization results",
+    )
+    parser.add_argument(
+        "-l",
+        "--labels",
+        default="./labels_coco_80.txt",
+        help="File to use for reading the class labels from, default: ./labels_coco_80.txt",
+    )
+    parser.add_argument(
+        "-c",
+        "--use_custom_gpu_allocator",
+        action="store_true",
+        default=False,
+        help="Use a custom gpu allocator with CUDA memory pools for better performance",
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", default=False, help="Set to verbose logging")
+    args = parser.parse_args()
+    main(args)

samples/python/dds_faster_rcnn/labels_coco_80.txt

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorcycle
+airplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+couch
+potted plant
+bed
+dining table
+toilet
+tv
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

samples/python/dds_faster_rcnn/modify_onnx.py

@@ -0,0 +1,56 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import onnx
+import numpy as np
+import argparse
+
+
+def modify_maskrcnn_opset12(path_to_model, output_path):
+    graph = gs.import_onnx(onnx.load(path_to_model))
+    """
+        Step 1: Remove unnecessary UINT8 cast
+            - Pattern match Cast[BOOL->UINT8] -> Cast[UINT8 -> BOOL]
+            - Fixes node 2838 - casts bool to uint8 for slice / gather. Can keep all operations in bool.
+    """
+    for node in graph.nodes:
+        if node.op == "Cast" and node.attrs["to"] == onnx.TensorProto.UINT8:
+            node.attrs["to"] = onnx.TensorProto.BOOL
+            node.outputs[0].dtype = np.bool_
+            # Need to modify output_node output to be bool as well.
+            for output_node in node.outputs[0].outputs:
+                output_node.outputs[0].dtype = np.bool_
+            print(f"Removed UINT8 casts in node {node.name}")
+
+    onnx.save(gs.export_onnx(graph.cleanup()), output_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-i",
+        "--input",
+        default="FasterRCNN-12.onnx",
+        help="Path to the onnx model obtained from https://github.com/onnx/models/raw/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-12.onnx",
+    )
+    parser.add_argument(
+        "-o", "--output", default="fasterrcnn12_trt.onnx", help="Desired path for the output onnx model"
+    )
+    args = parser.parse_args()
+
+    modify_maskrcnn_opset12(args.input, args.output)

samples/python/dds_faster_rcnn/requirements.txt

@@ -0,0 +1,5 @@
+Pillow==11.3.0
+cuda-python==12.9.0
+onnx==1.18.0
+onnx-graphsurgeon --index-url https://pypi.ngc.nvidia.com
+numpy==1.26.4

samples/python/dds_faster_rcnn/visualize.py

@@ -0,0 +1,199 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 numpy as np
+
+import PIL.Image as Image
+import PIL.ImageDraw as ImageDraw
+import PIL.ImageFont as ImageFont
+
+COLORS = [
+    "GoldenRod",
+    "MediumTurquoise",
+    "GreenYellow",
+    "SteelBlue",
+    "DarkSeaGreen",
+    "SeaShell",
+    "LightGrey",
+    "IndianRed",
+    "DarkKhaki",
+    "LawnGreen",
+    "WhiteSmoke",
+    "Peru",
+    "LightCoral",
+    "FireBrick",
+    "OldLace",
+    "LightBlue",
+    "SlateGray",
+    "OliveDrab",
+    "NavajoWhite",
+    "PaleVioletRed",
+    "SpringGreen",
+    "AliceBlue",
+    "Violet",
+    "DeepSkyBlue",
+    "Red",
+    "MediumVioletRed",
+    "PaleTurquoise",
+    "Tomato",
+    "Azure",
+    "Yellow",
+    "Cornsilk",
+    "Aquamarine",
+    "CadetBlue",
+    "CornflowerBlue",
+    "DodgerBlue",
+    "Olive",
+    "Orchid",
+    "LemonChiffon",
+    "Sienna",
+    "OrangeRed",
+    "Orange",
+    "DarkSalmon",
+    "Magenta",
+    "Wheat",
+    "Lime",
+    "GhostWhite",
+    "SlateBlue",
+    "Aqua",
+    "MediumAquaMarine",
+    "LightSlateGrey",
+    "MediumSeaGreen",
+    "SandyBrown",
+    "YellowGreen",
+    "Plum",
+    "FloralWhite",
+    "LightPink",
+    "Thistle",
+    "DarkViolet",
+    "Pink",
+    "Crimson",
+    "Chocolate",
+    "DarkGrey",
+    "Ivory",
+    "PaleGreen",
+    "DarkGoldenRod",
+    "LavenderBlush",
+    "SlateGrey",
+    "DeepPink",
+    "Gold",
+    "Cyan",
+    "LightSteelBlue",
+    "MediumPurple",
+    "ForestGreen",
+    "DarkOrange",
+    "Tan",
+    "Salmon",
+    "PaleGoldenRod",
+    "LightGreen",
+    "LightSlateGray",
+    "HoneyDew",
+    "Fuchsia",
+    "LightSeaGreen",
+    "DarkOrchid",
+    "Green",
+    "Chartreuse",
+    "LimeGreen",
+    "AntiqueWhite",
+    "Beige",
+    "Gainsboro",
+    "Bisque",
+    "SaddleBrown",
+    "Silver",
+    "Lavender",
+    "Teal",
+    "LightCyan",
+    "PapayaWhip",
+    "Purple",
+    "Coral",
+    "BurlyWood",
+    "LightGray",
+    "Snow",
+    "MistyRose",
+    "PowderBlue",
+    "DarkCyan",
+    "White",
+    "Turquoise",
+    "MediumSlateBlue",
+    "PeachPuff",
+    "Moccasin",
+    "LightSalmon",
+    "SkyBlue",
+    "Khaki",
+    "MediumSpringGreen",
+    "BlueViolet",
+    "MintCream",
+    "Linen",
+    "SeaGreen",
+    "HotPink",
+    "LightYellow",
+    "BlanchedAlmond",
+    "RoyalBlue",
+    "RosyBrown",
+    "MediumOrchid",
+    "DarkTurquoise",
+    "LightGoldenRodYellow",
+    "LightSkyBlue",
+]
+
+
+def visualize_detections(image_path, output_path, detections, labels=[]):
+    image = Image.open(image_path).convert(mode="RGB")
+    draw = ImageDraw.Draw(image)
+    line_width = 2
+    font = ImageFont.load_default()
+    for d in detections:
+        color = COLORS[d["class"] % len(COLORS)]
+        draw.line(
+            [
+                (d["xmin"], d["ymin"]),
+                (d["xmin"], d["ymax"]),
+                (d["xmax"], d["ymax"]),
+                (d["xmax"], d["ymin"]),
+                (d["xmin"], d["ymin"]),
+            ],
+            width=line_width,
+            fill=color,
+        )
+        label = f"Class {d['class']}"
+        if d["class"] < len(labels):
+            label = f"{labels[d['class']]}"
+        score = d["score"]
+        text = f"{label}: {int(100*score)}%"
+        if score < 0:
+            text = label
+        left, top, right, bottom = font.getbbox(text)
+        text_width, text_height = right - left, bottom - top
+        text_bottom = max(text_height, d["ymin"])
+        text_left = d["xmin"]
+        margin = np.ceil(0.05 * text_height)
+        draw.rectangle(
+            [
+                (text_left, text_bottom - text_height - 2 * margin),
+                (text_left + text_width, text_bottom),
+            ],
+            fill=color,
+        )
+        draw.text(
+            (text_left + margin, text_bottom - text_height - margin),
+            text,
+            fill="black",
+            font=font,
+        )
+    if output_path is None:
+        return image
+    image.save(output_path)