chore: import upstream snapshot with attribution

samples/python/detectron2/README.md

@@ -0,0 +1,191 @@
+# Detectron 2 Mask R-CNN R50-FPN 3x in TensorRT
+
+Support for Detectron 2 Mask R-CNN R50-FPN 3x model in TensorRT. This script helps with converting, running and validating this model with TensorRT.
+
+## Changelog
+
+- October 2025
+  - Migrate to strongly typed APIs.
+- Aug 2025
+  - Removed support for Python versions < 3.10.
+- Aug 2023
+  - Update ONNX version support to 1.14.0
+  - Update ONNX Runtime version support to 1.15.1 for Python>=3.8
+  - Removed support for Python versions < 3.8.
+- July 2023:
+  - Update benchmarks and include hardware used.
+- October 2022:
+  - Updated converter to support `tracing` export instead of deprecated `caffe2_tracing`
+
+## Setup
+
+In order for scripts to work we suggest an environment with TensorRT >= 8.4.1.
+
+Install TensorRT as per the [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.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.
+
+Install all dependencies listed in `requirements.txt`:
+
+```
+pip install -r requirements.txt
+```
+Note: this sample cannot be run on Jetson platforms as `torch.distributed` is not available. To check whether your platform supports `torch.distributed`, open a Python shell and confirm that `torch.distributed.is_available()` returns `True`.
+
+## Model Conversion
+
+The workflow to convert Detectron 2 Mask R-CNN R50-FPN 3x model is basically Detectron 2 → ONNX → TensorRT, and so parts of this process require Detectron 2 to be installed. Official export to ONNX is documented [here](https://detectron2.readthedocs.io/en/latest/tutorials/deployment.html).
+
+### Detectron 2 Deployment
+Deployment is done through export model script located in `detectron2/tools/deploy/export_model.py` of Detectron 2 [github](https://github.com/facebookresearch/detectron2). Detectron 2 Mask R-CNN R50-FPN 3x model is dynamic with minimum testing dimension size of 800 and maximum of 1333. TensorRT plug-ins used for conversion of this model do not support dynamic shapes, as a result we have to set both height and width of the input tensor to 1344. 1344 instead of 1333 because model requires both height and width of the input tensor to be divisible by 32. In order to export this model with correct 1344x1344 resolution, we have to make a change to `export_model.py`. Currently lines 160-162:
+
+```
+aug = T.ResizeShortestEdge(
+    [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
+)
+```
+have to be changed to:
+
+```
+aug = T.ResizeShortestEdge(
+    [1344, 1344], 1344
+)
+```
+
+Export script takes `--sample-image` as one of the arguments. Such image is used to adjust input dimensions and dimensions of tensors for the rest of the network. This sample image has to be an image of 1344x1344 dimensions, which contains at least one detectable by model object. My recommendation is to upsample one of COCO dataset images to 1344x1344. Sample command:
+
+```
+python detectron2/tools/deploy/export_model.py \
+    --sample-image 1344x1344.jpg \
+    --config-file detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --export-method tracing \
+    --format onnx \
+    --output ./ \
+    MODEL.WEIGHTS path/to/model_final_f10217.pkl \
+    MODEL.DEVICE cuda
+
+```
+
+Where `--sample-image` is 1344x1344 image; `--config-file` path to Mask R-CNN R50-FPN 3x config, included with detectron2; `MODEL.WEIGHTS` are weights of Mask R-CNN R50-FPN 3x that can be downloaded [here](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md). Resulted `model.onnx` will be an input to conversion script.
+
+### Create ONNX Graph
+This is supported Detectron 2 model:
+
+| **Model**                                     | **Resolution** |
+| ----------------------------------------------|----------------|
+| Mask R-CNN R50-FPN 3x                         | 1344x1344      |
+
+If Detectron 2 Mask R-CNN is ready to be converted (i.e. you ran `detectron2/tools/deploy/export_model.py`), run:
+
+```
+python create_onnx.py \
+    --exported_onnx /path/to/model.onnx \
+    --onnx /path/to/converted.onnx \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --det2_weights /model_final_f10217.pkl \
+    --sample_image any_image.jpg
+```
+
+This will create the file `converted.onnx` which is ready to convert to TensorRT.
+
+It is important to mention that `--sample_image` in this case is used for anchor generation. Detectron 2 ONNX models do not have anchor data inside the graph, so anchors have to be generated "offline". If custom model is used, make sure preprocessing of your model matches what is coded in `get_anchors(self, sample_image)` function.
+
+The script has a few optional arguments, including:
+
+* `--first_nms_threshold [...]` allows overriding the default 1st NMS score threshold parameter, as the runtime latency of the NMS plugin is sensitive to this value. It's a good practice to set this value as high as possible, while still fulfilling your application requirements, to reduce inference latency. In Mask R-CNN this will be a score threshold for Region Proposal Network.
+* `--second_nms_threshold [...]` allows overriding the default 2nd NMS score threshold parameter, further improves the runtime latency of the NMS plugin. It's a good practice to set this value as high as possible, while still fulfilling your application requirements, to reduce inference latency. It will be the second and last NMS.
+* `--batch_size` allows selection of various batch sizes, default is 1.
+
+
+Optionally, you may wish to visualize the resulting ONNX graph with a tool such as [Netron](https://netron.app/).
+
+The input to the graph is a `float32` tensor with the selected input shape, containing RGB pixel data in the range of 0 to 255. All preprocessing will be performed inside the Model graph, so it is not required to further pre-process the input data.
+
+
+The outputs of the graph are the same as the outputs of the [EfficientNMS_TRT](https://github.com/NVIDIA/TensorRT/tree/master/plugin/efficientNMSPlugin) plugin and segmentation head output, name of the last node is `detection_masks`, shape is `[batch_size, max_proposals, mask_height, mask_width]`, dtype is float32.
+
+### Build TensorRT Engine
+
+TensorRT engine can be built directly with `trtexec` using the ONNX graph generated in the previous step. If it's not already in your `$PATH`, the `trtexec` binary is usually found in `/usr/bin/trtexec`, depending on your TensorRT installation method. Run:
+
+```
+trtexec --onnx=/path/to/converted.onnx --saveEngine=/path/to/engine.trt --stronglyTyped
+```
+
+However, the script `build_engine.py` is also provided in this repository for convenience, as it has been tailored to Detectron 2 2 Mask R-CNN R50-FPN 3x engine building. Run `python3 build_engine.py --help` for details on available settings.
+
+#### Benchmark Engine
+
+Optionally, you can obtain execution timing information for the built engine by using the `trtexec` utility, as:
+
+```
+trtexec \
+    --loadEngine=/path/to/engine.trt \
+    --iterations=100 --avgRuns=100
+```
+
+An inference benchmark will run, with GPU Compute latency times printed out to the console. Depending on your environment, you should see something similar to:
+
+```
+GPU Compute Time: min = 30.1864 ms, max = 37.0945 ms, mean = 34.481 ms, median = 34.4187 ms, percentile(99%) = 37.0945 ms
+```
+
+Sample results with fp32 data precision are shown below. The following results were obtained using an RTX A5000 and TensorRT 8.6.1. mAP was evaluated for the COCO val2017 dataset using the instructions in [Evaluate mAP Metric](#evaluate-map-metric).
+
+| **Precision**   | **Latency** | **bbox COCO mAP** | **segm COCO mAP** |
+| ----------------|-------------|-------------------|-------------------|
+| fp32            | 25.89 ms    | 0.402             | 0.368             |
+
+## Inference
+
+For optimal performance, inference should be done in a C++ application that takes advantage of CUDA Graphs to launch the inference request. Alternatively, the TensorRT engine built with this process can also be executed through either [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server) or [DeepStream SDK](https://developer.nvidia.com/deepstream-sdk).
+
+However, for convenience, a python inference script is provided here for quick testing of the built TensorRT engine.
+
+### Inference in Python
+
+To perform object detection on a set of images with TensorRT, run:
+
+```
+python infer.py \
+    --engine /path/to/engine.trt \
+    --input /path/to/images \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --output /path/to/output \
+```
+
+Where the input path can be either a single image file, or a directory of jpg/png/bmp images.
+
+The script has a few optional arguments, including:
+* `--nms_threshold` allows overriding the default second NMS score threshold parameter.
+* `--iou_threshold` allows to set IoU threshold for the mask segmentation, default is 0.5.
+
+The detection results will be written out to the specified output directory, consisting of a visualization image, and a tab-separated results file for each input image processed.
+
+#### Sample Images
+
+![infer_1](https://drive.google.com/uc?export=view&id=1AOW9IXqjrU7eVYmaue-pqijNucXmx_s0)
+
+![infer_2](https://drive.google.com/uc?export=view&id=1m1fp2v41DOqKfj423G0-eyKVurrPNYGx)
+
+### Evaluate mAP Metric
+
+Given a validation dataset (such as [COCO val2017 data](http://images.cocodataset.org/zips/val2017.zip)), you can get the mAP metrics for the built TensorRT engine. This will use the mAP metrics tools functions from the [Detectron 2 evaluation](https://github.com/facebookresearch/detectron2/tree/main/detectron2/evaluation) repository. Make sure you follow [Use Builtin Datasets guide](https://detectron2.readthedocs.io/en/latest/tutorials/builtin_datasets.html) to correctly setup COCO or custom dataset. Additionally, run `eval_coco.py` in the same folder where `/datasets` is present, otherwise this error will appear:
+
+```
+FileNotFoundError: [Errno 2] No such file or directory: 'datasets/coco/annotations/instances_val2017.json'
+```
+
+To run evalutions, run:
+
+```
+python eval_coco.py \
+    --engine /path/to/engine.trt \
+    --input /path/to/coco/val2017 \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --det2_weights /model_final_f10217.pkl
+```
+
+The script has a few optional arguments, including:
+* `--nms_threshold` allows overriding the default second NMS score threshold parameter.
+* `--iou_threshold` allows to set IoU threshold for the mask segmentation, default is 0.5.
+
+The mAP metric is sensitive to the NMS score threshold used, as using a high threshold will reduce the model recall, resulting in a lower mAP value. It may be a good idea to build separate TensorRT engines for different purposes. That is, one engine with a default threshold (like 0.5) dedicated for mAP validation, and another engine with your application specific threshold (like 0.8) for deployment. This is why we keep the NMS threshold as a configurable parameter in the `create_onnx.py` script.

samples/python/detectron2/build_engine.py

@@ -0,0 +1,147 @@
+#
+# 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 numpy as np
+import tensorrt as trt
+from cuda.bindings import driver as cuda, runtime as cudart
+
+from image_batcher import ImageBatcher
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("EngineBuilder").setLevel(logging.INFO)
+log = logging.getLogger("EngineBuilder")
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+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()
+        self.config.set_memory_pool_limit(
+            trt.MemoryPoolType.WORKSPACE, workspace * (2**30)
+        )
+
+        self.batch_size = None
+        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()):
+                log.error("Failed to load ONNX file: {}".format(onnx_path))
+                for error in range(self.parser.num_errors):
+                    log.error(self.parser.get_error(error))
+                sys.exit(1)
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+        outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)]
+
+        log.info("Network Description")
+        for input in inputs:
+            self.batch_size = input.shape[0]
+            log.info(
+                "Input '{}' with shape {} and dtype {}".format(
+                    input.name, input.shape, input.dtype
+                )
+            )
+        for output in outputs:
+            log.info(
+                "Output '{}' with shape {} and dtype {}".format(
+                    output.name, output.shape, output.dtype
+                )
+            )
+        assert self.batch_size > 0
+
+    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("Building fp32 Engine in {}".format(engine_path))
+
+        engine_bytes = self.builder.build_serialized_network(self.network, self.config)
+        if engine_bytes is None:
+            log.error("Failed to create engine")
+            sys.exit(1)
+
+        with open(engine_path, "wb") as f:
+            log.info("Serializing engine to file: {:}".format(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", help="The input ONNX model file to load")
+    parser.add_argument("-e", "--engine", 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=1,
+        type=int,
+        help="The max memory workspace size to allow in Gb, " "default: 1",
+    )
+    args = parser.parse_args()
+    if not all([args.onnx, args.engine]):
+        parser.print_help()
+        log.error("These arguments are required: --onnx and --engine")
+        sys.exit(1)
+    main(args)

samples/python/detectron2/create_onnx.py

@@ -0,0 +1,879 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 re
+import sys
+import argparse
+import logging
+import cv2
+import onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+from onnx import shape_inference
+import torch
+
+try:
+    from detectron2.engine.defaults import DefaultPredictor
+    from detectron2.modeling import build_model
+    from detectron2.config import get_cfg
+    from detectron2.structures import ImageList
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+import onnx_utils
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("ModelHelper").setLevel(logging.INFO)
+log = logging.getLogger("ModelHelper")
+
+
+class DET2GraphSurgeon:
+    def __init__(self, saved_model_path, config_file, weights):
+        """
+        Constructor of the Model Graph Surgeon object, to do the conversion of a Detectron 2 Mask R-CNN exported model
+        to an ONNX-TensorRT parsable model.
+        :param saved_model_path: The path pointing to the exported Detectron 2 Mask R-CNN ONNX model.
+        :param config_file: The path pointing to the Detectron 2 yaml file which describes the model.
+        :param config_file: Weights to load for the Detectron 2 model.
+        """
+
+        def det2_setup(config_file, weights):
+            """
+            Create configs and perform basic setups.
+            """
+            cfg = get_cfg()
+            cfg.merge_from_file(config_file)
+            cfg.merge_from_list(["MODEL.WEIGHTS", weights])
+            cfg.freeze()
+            return cfg
+
+        # Import exported Detectron 2 Mask R-CNN ONNX model as GraphSurgeon object.
+        self.graph = gs.import_onnx(onnx.load(saved_model_path))
+        assert self.graph
+        log.info("ONNX graph loaded successfully")
+
+        # Fold constants via ONNX-GS that exported script might've missed.
+        self.graph.fold_constants()
+
+        # Set up Detectron 2 model configuration.
+        self.det2_cfg = det2_setup(config_file, weights)
+
+        # Getting model characteristics.
+        self.fpn_out_channels = self.det2_cfg.MODEL.FPN.OUT_CHANNELS
+        self.num_classes = self.det2_cfg.MODEL.ROI_HEADS.NUM_CLASSES
+        self.first_NMS_max_proposals = self.det2_cfg.MODEL.RPN.POST_NMS_TOPK_TEST
+        self.first_NMS_iou_threshold = self.det2_cfg.MODEL.RPN.NMS_THRESH
+        self.first_NMS_score_threshold = 0.01
+        self.first_ROIAlign_pooled_size = (
+            self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION
+        )
+        self.first_ROIAlign_sampling_ratio = (
+            self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
+        )
+        self.first_ROIAlign_type = self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE
+        self.second_NMS_max_proposals = self.det2_cfg.TEST.DETECTIONS_PER_IMAGE
+        self.second_NMS_iou_threshold = self.det2_cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST
+        self.second_NMS_score_threshold = (
+            self.det2_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST
+        )
+        self.second_ROIAlign_pooled_size = (
+            self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION
+        )
+        self.second_ROIAlign_sampling_ratio = (
+            self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_SAMPLING_RATIO
+        )
+        self.second_ROIAlign_type = self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_TYPE
+        self.mask_out_res = 28
+
+        # Model characteristics.
+        log.info("Number of FPN output channels is {}".format(self.fpn_out_channels))
+        log.info("Number of classes is {}".format(self.num_classes))
+        log.info("First NMS max proposals is {}".format(self.first_NMS_max_proposals))
+        log.info("First NMS iou threshold is {}".format(self.first_NMS_iou_threshold))
+        log.info(
+            "First NMS score threshold is {}".format(self.first_NMS_score_threshold)
+        )
+        log.info("First ROIAlign type is {}".format(self.first_ROIAlign_type))
+        log.info(
+            "First ROIAlign pooled size is {}".format(self.first_ROIAlign_pooled_size)
+        )
+        log.info(
+            "First ROIAlign sampling ratio is {}".format(
+                self.first_ROIAlign_sampling_ratio
+            )
+        )
+        log.info("Second NMS max proposals is {}".format(self.second_NMS_max_proposals))
+        log.info("Second NMS iou threshold is {}".format(self.second_NMS_iou_threshold))
+        log.info(
+            "Second NMS score threshold is {}".format(self.second_NMS_score_threshold)
+        )
+        log.info("Second ROIAlign type is {}".format(self.second_ROIAlign_type))
+        log.info(
+            "Second ROIAlign pooled size is {}".format(self.second_ROIAlign_pooled_size)
+        )
+        log.info(
+            "Second ROIAlign sampling ratio is {}".format(
+                self.second_ROIAlign_sampling_ratio
+            )
+        )
+        log.info(
+            "Individual mask output resolution is {}x{}".format(
+                self.mask_out_res, self.mask_out_res
+            )
+        )
+
+        self.batch_size = None
+
+    def sanitize(self):
+        """
+        Sanitize the graph by cleaning any unconnected nodes, do a topological resort, and fold constant inputs values.
+        When possible, run shape inference on the ONNX graph to determine tensor shapes.
+        """
+
+        for i in range(3):
+            count_before = len(self.graph.nodes)
+            self.graph.cleanup().toposort()
+            try:
+                for node in self.graph.nodes:
+                    for o in node.outputs:
+                        o.shape = None
+                model = gs.export_onnx(self.graph)
+                model = shape_inference.infer_shapes(model)
+                self.graph = gs.import_onnx(model)
+            except Exception as e:
+                log.info(
+                    "Shape inference could not be performed at this time:\n{}".format(e)
+                )
+            try:
+                self.graph.fold_constants(fold_shapes=True)
+            except TypeError as e:
+                log.error(
+                    "This version of ONNX GraphSurgeon does not support folding shapes, please upgrade your "
+                    "onnx_graphsurgeon module. Error:\n{}".format(e)
+                )
+                raise
+
+            count_after = len(self.graph.nodes)
+            if count_before == count_after:
+                # No new folding occurred in this iteration, so we can stop for now.
+                break
+
+    def get_anchors(self, sample_image):
+        """
+        Detectron 2 exported ONNX does not contain anchors required for efficientNMS plug-in, so they must be generated
+        "offline" by calling actual Detectron 2 model and getting anchors from it.
+        :param sample_image: Sample image required to run through the model and obtain anchors.
+        Can be any image from a dataset. Make sure listed here Detectron 2 preprocessing steps
+        actually match your preprocessing steps. Otherwise, behavior can be unpredictable.
+        Additionally, anchors have to be generated for a fixed input dimensions,
+        meaning as soon as image leaves a preprocessor and enters predictor.model.backbone() it must have
+        a fixed dimension (1344x1344 in my case) that every single image in dataset must follow, since currently
+        TensorRT plug-ins do not support dynamic shapes.
+        """
+        # Get Detectron 2 model config and build it.
+        predictor = DefaultPredictor(self.det2_cfg)
+        model = build_model(self.det2_cfg)
+
+        # Image preprocessing.
+        input_im = cv2.imread(sample_image)
+        raw_height, raw_width = input_im.shape[:2]
+        image = predictor.aug.get_transform(input_im).apply_image(input_im)
+        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+
+        # Model preprocessing.
+        inputs = [{"image": image, "height": raw_height, "width": raw_width}]
+        images = [x["image"].to(model.device) for x in inputs]
+        images = [(x - model.pixel_mean) / model.pixel_std for x in images]
+        imagelist_images = ImageList.from_tensors(images, 1344)
+
+        # Get feature maps from backbone.
+        features = predictor.model.backbone(imagelist_images.tensor)
+
+        # Get proposals from Region Proposal Network and obtain anchors from anchor generator.
+        features = [features[f] for f in predictor.model.proposal_generator.in_features]
+        det2_anchors = predictor.model.proposal_generator.anchor_generator(features)
+
+        # Extract anchors based on feature maps in ascending order (P2->P6).
+        p2_anchors = det2_anchors[0].tensor.detach().cpu().numpy()
+        p3_anchors = det2_anchors[1].tensor.detach().cpu().numpy()
+        p4_anchors = det2_anchors[2].tensor.detach().cpu().numpy()
+        p5_anchors = det2_anchors[3].tensor.detach().cpu().numpy()
+        p6_anchors = det2_anchors[4].tensor.detach().cpu().numpy()
+        final_anchors = np.concatenate(
+            (p2_anchors, p3_anchors, p4_anchors, p5_anchors, p6_anchors)
+        )
+
+        return final_anchors
+
+    def save(self, output_path):
+        """
+        Save the ONNX model to the given location.
+        :param output_path: Path pointing to the location where to write out the updated ONNX model.
+        """
+        self.graph.cleanup().toposort()
+        model = gs.export_onnx(self.graph)
+        output_path = os.path.realpath(output_path)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        onnx.save(model, output_path)
+        log.info("Saved ONNX model to {}".format(output_path))
+
+    def update_preprocessor(self, batch_size):
+        """
+        Remove all the pre-processing nodes in the ONNX graph and leave only the image normalization essentials.
+        :param batch_size: The batch size to use for the ONNX graph.
+        """
+        # Set graph inputs.
+        self.batch_size = batch_size
+        self.height = self.graph.inputs[0].shape[1]
+        self.width = self.graph.inputs[0].shape[2]
+
+        input_shape = [self.batch_size, 3, self.height, self.width]
+        self.graph.inputs[0].shape = input_shape
+        self.graph.inputs[0].dtype = np.float32
+        self.graph.inputs[0].name = "input_tensor"
+
+        self.sanitize()
+        log.info(
+            "ONNX graph input shape: {} [NCHW format set]".format(
+                self.graph.inputs[0].shape
+            )
+        )
+
+        # Find the initial nodes of the graph, whatever the input is first connected to, and disconnect them.
+        for node in [
+            node for node in self.graph.nodes if self.graph.inputs[0] in node.inputs
+        ]:
+            node.inputs.clear()
+
+        # Get input tensor.
+        input_tensor = self.graph.inputs[0]
+
+        # Create preprocessing Sub node and connect input tensor to it.
+        sub_const = np.expand_dims(
+            np.asarray([255 * 0.406, 255 * 0.456, 255 * 0.485], dtype=np.float32),
+            axis=(1, 2),
+        )
+        sub_out = self.graph.op_with_const(
+            "Sub", "preprocessor/mean", input_tensor, sub_const
+        )
+
+        # Find first Div node and connect to output of Sub node.
+        div_node = self.graph.find_node_by_op("Div")
+        log.info("Found {} node".format(div_node.op))
+        div_node.inputs[0] = sub_out[0]
+
+        # Find first Conv and connect preprocessor directly to it.
+        conv_node = self.graph.find_node_by_op("Conv")
+        log.info("Found {} node".format(conv_node.op))
+        conv_node.inputs[0] = div_node.outputs[0]
+
+        # Reshape nodes tend to update the batch dimension to a fixed value of 1, they should use the batch size instead.
+        for node in [node for node in self.graph.nodes if node.op == "Reshape"]:
+            if type(node.inputs[1]) == gs.Constant and node.inputs[1].values[0] == 1:
+                node.inputs[1].values[0] = self.batch_size
+
+    def NMS(
+        self,
+        boxes,
+        scores,
+        anchors,
+        background_class,
+        score_activation,
+        max_proposals,
+        iou_threshold,
+        nms_score_threshold,
+        user_threshold,
+        nms_name=None,
+    ):
+        # Helper function to create the NMS Plugin node with the selected inputs.
+        # EfficientNMS_TRT TensorRT Plugin is suitable for our use case.
+        # :param boxes: The box predictions from the Box Net.
+        # :param scores: The class predictions from the Class Net.
+        # :param anchors: The default anchor coordinates.
+        # :param background_class: The label ID for the background class.
+        # :param max_proposals: Number of proposals made by NMS.
+        # :param score_activation: If set to True - apply sigmoid activation to the confidence scores during NMS operation,
+        # if false - no activation.
+        # :param iou_threshold: NMS intersection over union threshold, given by self.det2_cfg.
+        # :param nms_score_threshold: NMS score threshold, given by self.det2_cfg.
+        # :param user_threshold: User's given threshold to overwrite default NMS score threshold.
+        # :param nms_name: Name of NMS node in a graph, renames NMS elements accordingly in order to eliminate cycles.
+
+        if nms_name is None:
+            nms_name = ""
+        else:
+            nms_name = "_" + nms_name
+
+        # Set score threshold.
+        score_threshold = (
+            nms_score_threshold if user_threshold is None else user_threshold
+        )
+
+        # NMS Outputs.
+        nms_output_num_detections = gs.Variable(
+            name="num_detections" + nms_name, dtype=np.int32, shape=[self.batch_size, 1]
+        )
+        nms_output_boxes = gs.Variable(
+            name="detection_boxes" + nms_name,
+            dtype=np.float32,
+            shape=[self.batch_size, max_proposals, 4],
+        )
+        nms_output_scores = gs.Variable(
+            name="detection_scores" + nms_name,
+            dtype=np.float32,
+            shape=[self.batch_size, max_proposals],
+        )
+        nms_output_classes = gs.Variable(
+            name="detection_classes" + nms_name,
+            dtype=np.int32,
+            shape=[self.batch_size, max_proposals],
+        )
+
+        nms_outputs = [
+            nms_output_num_detections,
+            nms_output_boxes,
+            nms_output_scores,
+            nms_output_classes,
+        ]
+
+        # Plugin.
+        self.graph.plugin(
+            op="EfficientNMS_TRT",
+            name="nms" + nms_name,
+            inputs=[boxes, scores, anchors],
+            outputs=nms_outputs,
+            attrs={
+                "plugin_version": "1",
+                "background_class": background_class,
+                "max_output_boxes": max_proposals,
+                "score_threshold": max(0.01, score_threshold),
+                "iou_threshold": iou_threshold,
+                "score_activation": score_activation,
+                "class_agnostic": False,
+                "box_coding": 1,
+            },
+        )
+        log.info("Created nms{} with EfficientNMS_TRT plugin".format(nms_name))
+
+        return nms_outputs
+
+    def ROIAlign(
+        self,
+        rois,
+        p2,
+        p3,
+        p4,
+        p5,
+        pooled_size,
+        sampling_ratio,
+        roi_align_type,
+        num_rois,
+        ra_name,
+    ):
+        # Helper function to create the ROIAlign Plugin node with the selected inputs.
+        # PyramidROIAlign_TRT TensorRT Plugin is suitable for our use case.
+        # :param rois: Regions of interest/detection boxes outputs from preceding NMS node.
+        # :param p2: Output of p2 feature map.
+        # :param p3: Output of p3 feature map.
+        # :param p4: Output of p4 feature map.
+        # :param p5: Output of p5 feature map.
+        # :param pooled_size: Pooled output dimensions.
+        # :param sampling_ratio: Number of sampling points in the interpolation grid used to compute the output value of each pooled output bin.
+        # :param roi_align_type: Type of Detectron 2 ROIAlign op, either ROIAlign (vanilla) or ROIAlignV2 (0.5 coordinate offset).
+        # :param num_rois: Number of ROIs resulting from ROIAlign operation.
+        # :param ra_name: Name of ROIAlign node in a graph, renames ROIAlign elements accordingly in order to eliminate cycles.
+
+        # Different types of Detectron 2's ROIAlign ops require coordinate offset that is supported by PyramidROIAlign_TRT.
+        if roi_align_type == "ROIAlignV2":
+            roi_coords_transform = 2
+        elif roi_align_type == "ROIAlign":
+            roi_coords_transform = 0
+        else:
+            raise ValueError(
+                f"Unsupported ROIAlign type '{roi_align_type}'. "
+                f"Expected 'ROIAlignV2' or 'ROIAlign'."
+            )
+
+        # ROIAlign outputs.
+        roi_align_output = gs.Variable(
+            name="roi_align/output_" + ra_name,
+            dtype=np.float32,
+            shape=[
+                self.batch_size,
+                num_rois,
+                self.fpn_out_channels,
+                pooled_size,
+                pooled_size,
+            ],
+        )
+
+        # Plugin.
+        self.graph.plugin(
+            op="PyramidROIAlign_TRT",
+            name="roi_align_" + ra_name,
+            inputs=[rois, p2, p3, p4, p5],
+            outputs=[roi_align_output],
+            attrs={
+                "plugin_version": "1",
+                "fpn_scale": 224,
+                "pooled_size": pooled_size,
+                "image_size": [self.height, self.width],
+                "roi_coords_absolute": 0,
+                "roi_coords_swap": 0,
+                "roi_coords_transform": roi_coords_transform,
+                "sampling_ratio": sampling_ratio,
+            },
+        )
+        log.info("Created {} with PyramidROIAlign_TRT plugin".format(ra_name))
+
+        return roi_align_output
+
+    def process_graph(
+        self, anchors, first_nms_threshold=None, second_nms_threshold=None
+    ):
+        """
+        Processes the graph to replace the GenerateProposals and BoxWithNMSLimit operations with EfficientNMS_TRT
+        TensorRT plugin nodes and ROIAlign operations with PyramidROIAlign_TRT plugin nodes.
+        :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided
+        inside the graph.
+        :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph.
+        :param second_nms_threshold: Override the 2nd NMS score threshold value. If set to None, use the value in the graph.
+        """
+
+        def backbone():
+            """
+            Updates the graph to replace all ResizeNearest ops with ResizeNearest plugins in backbone.
+            """
+            # Get final backbone outputs.
+            p2 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output2/Conv")
+            p3 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output3/Conv")
+            p4 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output4/Conv")
+            p5 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output5/Conv")
+
+            return p2.outputs[0], p3.outputs[0], p4.outputs[0], p5.outputs[0]
+
+        def proposal_generator(anchors, first_nms_threshold):
+            """
+            Updates the graph to replace all GenerateProposals Caffe ops with one single NMS for proposals generation.
+            :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided
+            inside the graph
+            :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph.
+            """
+            # Get nodes containing final objectness logits.
+            p2_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten"
+            )
+            p3_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_1"
+            )
+            p4_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_2"
+            )
+            p5_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_3"
+            )
+            p6_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_4"
+            )
+
+            # Get nodes containing final anchor_deltas.
+            p2_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_1"
+            )
+            p3_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_3"
+            )
+            p4_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_5"
+            )
+            p5_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_7"
+            )
+            p6_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_9"
+            )
+
+            # Concatenate all objectness logits/scores data.
+            scores_inputs = [
+                p2_logits.outputs[0],
+                p3_logits.outputs[0],
+                p4_logits.outputs[0],
+                p5_logits.outputs[0],
+                p6_logits.outputs[0],
+            ]
+            scores_tensor = self.graph.layer(
+                name="scores",
+                op="Concat",
+                inputs=scores_inputs,
+                outputs=["scores"],
+                attrs={"axis": 1},
+            )[0]
+            # Unsqueeze to add 3rd dimension of 1 to match tensor dimensions of boxes tensor.
+            scores = self.graph.unsqueeze("scores_unsqueeze", scores_tensor, [2])[0]
+
+            # Concatenate all boxes/anchor_delta data.
+            boxes_inputs = [
+                p2_anchors.outputs[0],
+                p3_anchors.outputs[0],
+                p4_anchors.outputs[0],
+                p5_anchors.outputs[0],
+                p6_anchors.outputs[0],
+            ]
+            boxes = self.graph.layer(
+                name="boxes",
+                op="Concat",
+                inputs=boxes_inputs,
+                outputs=["anchors"],
+                attrs={"axis": 1},
+            )[0]
+
+            # Convert the anchors from Corners to CenterSize encoding.
+            anchors = np.matmul(
+                anchors,
+                [[0.5, 0, -1, 0], [0, 0.5, 0, -1], [0.5, 0, 1, 0], [0, 0.5, 0, 1]],
+            )
+            anchors = anchors / [
+                self.width,
+                self.height,
+                self.width,
+                self.height,
+            ]  # Normalize anchors to [0-1] range
+            anchors = np.expand_dims(anchors, axis=0)
+            anchors = anchors.astype(np.float32)
+            anchors = gs.Constant(name="default_anchors", values=anchors)
+
+            # Create NMS node.
+            nms_outputs = self.NMS(
+                boxes,
+                scores,
+                anchors,
+                -1,
+                False,
+                self.first_NMS_max_proposals,
+                self.first_NMS_iou_threshold,
+                self.first_NMS_score_threshold,
+                first_nms_threshold,
+                "rpn",
+            )
+
+            return nms_outputs
+
+        def roi_heads(rpn_outputs, p2, p3, p4, p5, second_nms_threshold):
+            """
+            Updates the graph to replace all ROIAlign Caffe ops with one single pyramid ROIAlign. Eliminates CollectRpnProposals
+            DistributeFpnProposals and BatchPermutation nodes that are not supported by TensorRT. Connects pyramid ROIAlign to box_head
+            and connects box_head to final box head outputs in a form of second NMS. In order to implement mask head outputs,
+            similar steps as in box_pooler are performed to replace mask_pooler. Finally, reimplemented mask_pooler is connected to
+            mask_head and mask head outputs are produced.
+            :param rpn_outputs: Outputs of the first NMS/proposal generator.
+            :param p2: Output of p2 feature map, required for ROIAlign operation.
+            :param p3: Output of p3 feature map, required for ROIAlign operation.
+            :param p4: Output of p4 feature map, required for ROIAlign operation.
+            :param p5: Output of p5 feature map, required for ROIAlign operation.
+            :param second_nms_threshold: Override the 2nd NMS score threshold value. If set to None, use the value in the graph.
+            """
+            # Create ROIAlign node.
+            box_pooler_output = self.ROIAlign(
+                rpn_outputs[1],
+                p2,
+                p3,
+                p4,
+                p5,
+                self.first_ROIAlign_pooled_size,
+                self.first_ROIAlign_sampling_ratio,
+                self.first_ROIAlign_type,
+                self.first_NMS_max_proposals,
+                "box_pooler",
+            )
+
+            # Reshape node that prepares ROIAlign/box pooler output for Gemm node that comes next.
+            box_pooler_shape = np.asarray(
+                [
+                    -1,
+                    self.fpn_out_channels
+                    * self.first_ROIAlign_pooled_size
+                    * self.first_ROIAlign_pooled_size,
+                ],
+                dtype=np.int64,
+            )
+            box_pooler_reshape = self.graph.op_with_const(
+                "Reshape", "box_pooler/reshape", box_pooler_output, box_pooler_shape
+            )
+
+            # Get first Gemm op of box head and connect box pooler to it.
+            first_box_head_gemm = self.graph.find_node_by_op_name(
+                "Gemm", "/roi_heads/box_head/fc1/Gemm"
+            )
+            first_box_head_gemm.inputs[0] = box_pooler_reshape[0]
+
+            # Get final two nodes of box predictor. Softmax op for cls_score, Gemm op for bbox_pred.
+            cls_score = self.graph.find_node_by_op_name("Softmax", "/roi_heads/Softmax")
+            bbox_pred = self.graph.find_node_by_op_name(
+                "Gemm", "/roi_heads/box_predictor/bbox_pred/Gemm"
+            )
+
+            # Linear transformation to convert box coordinates from (TopLeft, BottomRight) Corner encoding
+            # to CenterSize encoding. 1st NMS boxes are multiplied by transformation matrix in order to
+            # encode it into CenterSize format.
+            matmul_const = np.matrix(
+                "0.5 0 -1 0; 0 0.5 0 -1; 0.5 0 1 0; 0 0.5 0 1", dtype=np.float32
+            )
+            matmul_out = self.graph.matmul(
+                "RPN_NMS/detection_boxes_conversion", rpn_outputs[1], matmul_const
+            )
+
+            # Reshape node that prepares bbox_pred for scaling and second NMS.
+            bbox_pred_shape = np.asarray(
+                [self.batch_size, self.first_NMS_max_proposals, self.num_classes, 4],
+                dtype=np.int64,
+            )
+            bbox_pred_reshape = self.graph.op_with_const(
+                "Reshape", "bbox_pred/reshape", bbox_pred.outputs[0], bbox_pred_shape
+            )
+
+            # 0.1, 0.1, 0.2, 0.2 are localization head variance numbers, they scale bbox_pred_reshape, in order to get accurate coordinates.
+            scale_adj = np.expand_dims(
+                np.asarray([0.1, 0.1, 0.2, 0.2], dtype=np.float32), axis=(0, 1)
+            )
+            final_bbox_pred = self.graph.op_with_const(
+                "Mul", "bbox_pred/scale", bbox_pred_reshape[0], scale_adj
+            )
+
+            # Reshape node that prepares cls_score for slicing and second NMS.
+            cls_score_shape = np.array(
+                [self.batch_size, self.first_NMS_max_proposals, self.num_classes + 1],
+                dtype=np.int64,
+            )
+            cls_score_reshape = self.graph.op_with_const(
+                "Reshape", "cls_score/reshape", cls_score.outputs[0], cls_score_shape
+            )
+
+            # Slice operation to adjust third dimension of cls_score tensor, deletion of background class (81 in Detectron 2).
+            final_cls_score = self.graph.slice(
+                "cls_score/slicer", cls_score_reshape[0], 0, self.num_classes, 2
+            )
+
+            # Create NMS node.
+            nms_outputs = self.NMS(
+                final_bbox_pred[0],
+                final_cls_score[0],
+                matmul_out[0],
+                -1,
+                False,
+                self.second_NMS_max_proposals,
+                self.second_NMS_iou_threshold,
+                self.second_NMS_score_threshold,
+                second_nms_threshold,
+                "box_outputs",
+            )
+
+            # Create ROIAlign node.
+            mask_pooler_output = self.ROIAlign(
+                nms_outputs[1],
+                p2,
+                p3,
+                p4,
+                p5,
+                self.second_ROIAlign_pooled_size,
+                self.second_ROIAlign_sampling_ratio,
+                self.second_ROIAlign_type,
+                self.second_NMS_max_proposals,
+                "mask_pooler",
+            )
+
+            # Reshape mask pooler output.
+            mask_pooler_shape = np.asarray(
+                [
+                    self.second_NMS_max_proposals * self.batch_size,
+                    self.fpn_out_channels,
+                    self.second_ROIAlign_pooled_size,
+                    self.second_ROIAlign_pooled_size,
+                ],
+                dtype=np.int64,
+            )
+            mask_pooler_reshape_node = self.graph.op_with_const(
+                "Reshape", "mask_pooler/reshape", mask_pooler_output, mask_pooler_shape
+            )
+
+            # Get first Conv op in mask head and connect ROIAlign's squeezed output to it.
+            mask_head_conv = self.graph.find_node_by_op_name(
+                "Conv", "/roi_heads/mask_head/mask_fcn1/Conv"
+            )
+            mask_head_conv.inputs[0] = mask_pooler_reshape_node[0]
+
+            # Reshape node that is preparing 2nd NMS class outputs for Add node that comes next.
+            classes_reshape_shape = np.asarray(
+                [self.second_NMS_max_proposals * self.batch_size], dtype=np.int64
+            )
+            classes_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "box_outputs/reshape_classes",
+                nms_outputs[3],
+                classes_reshape_shape,
+            )
+
+            # This loop will generate an array used in Add node, which eventually will help Gather node to pick the single
+            # class of interest per bounding box, instead of creating 80 masks for every single bounding box.
+            add_array = []
+            for i in range(self.second_NMS_max_proposals * self.batch_size):
+                if i == 0:
+                    start_pos = 0
+                else:
+                    start_pos = i * self.num_classes
+                add_array.append(start_pos)
+
+            # This Add node is one of the Gather node inputs, Gather node performs gather on 0th axis of data tensor
+            # and requires indices that set tensors to be withing bounds, this Add node provides the bounds for Gather.
+            add_array = np.asarray(add_array, dtype=np.int32)
+            classes_add_node = self.graph.op_with_const(
+                "Add", "box_outputs/add", classes_reshape_node[0], add_array
+            )
+
+            # Get the last Conv op in mask head and reshape it to correctly gather class of interest's masks.
+            last_conv = self.graph.find_node_by_op_name(
+                "Conv", "/roi_heads/mask_head/predictor/Conv"
+            )
+            last_conv_reshape_shape = np.asarray(
+                [
+                    self.second_NMS_max_proposals * self.num_classes * self.batch_size,
+                    self.mask_out_res,
+                    self.mask_out_res,
+                ],
+                dtype=np.int64,
+            )
+            last_conv_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "mask_head/reshape_all_masks",
+                last_conv.outputs[0],
+                last_conv_reshape_shape,
+            )
+
+            # Gather node that selects only masks belonging to detected class, 79 other masks are discarded.
+            final_gather = self.graph.gather(
+                "mask_head/final_gather",
+                last_conv_reshape_node[0],
+                classes_add_node[0],
+                0,
+            )
+
+            # Get last Sigmoid node and connect Gather node to it.
+            mask_head_sigmoid = self.graph.find_node_by_op_name(
+                "Sigmoid", "/roi_heads/mask_head/Sigmoid"
+            )
+            mask_head_sigmoid.inputs[0] = final_gather[0]
+
+            # Final Reshape node, reshapes output of Sigmoid, important for various batch_size support (not tested yet).
+            final_graph_reshape_shape = np.asarray(
+                [
+                    self.batch_size,
+                    self.second_NMS_max_proposals,
+                    self.mask_out_res,
+                    self.mask_out_res,
+                ],
+                dtype=np.int64,
+            )
+            final_graph_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "mask_head/final_reshape",
+                mask_head_sigmoid.outputs[0],
+                final_graph_reshape_shape,
+            )
+            final_graph_reshape_node[0].dtype = np.float32
+            final_graph_reshape_node[0].name = "detection_masks"
+
+            return nms_outputs, final_graph_reshape_node[0]
+
+        # Only Detectron 2's Mask-RCNN R50-FPN 3x is supported currently.
+        p2, p3, p4, p5 = backbone()
+        rpn_outputs = proposal_generator(anchors, first_nms_threshold)
+        box_head_outputs, mask_head_output = roi_heads(
+            rpn_outputs, p2, p3, p4, p5, second_nms_threshold
+        )
+        # Append segmentation head output.
+        box_head_outputs.append(mask_head_output)
+        # Set graph outputs, both bbox and segmentation heads.
+        self.graph.outputs = box_head_outputs
+        self.sanitize()
+
+
+def main(args):
+    det2_gs = DET2GraphSurgeon(args.exported_onnx, args.det2_config, args.det2_weights)
+    det2_gs.update_preprocessor(args.batch_size)
+    anchors = det2_gs.get_anchors(args.sample_image)
+    det2_gs.process_graph(anchors, args.first_nms_threshold, args.second_nms_threshold)
+    det2_gs.save(args.onnx)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-i",
+        "--exported_onnx",
+        help="The exported to ONNX Detectron 2 Mask R-CNN",
+        type=str,
+    )
+    parser.add_argument(
+        "-o", "--onnx", help="The output ONNX model file to write", type=str
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-w", "--det2_weights", help="The Detectron 2 model weights (.pkl)", type=str
+    )
+    parser.add_argument(
+        "-s", "--sample_image", help="Sample image for anchors generation", type=str
+    )
+    parser.add_argument(
+        "-b", "--batch_size", help="Batch size for the model", type=int, default=1
+    )
+    parser.add_argument(
+        "-t1",
+        "--first_nms_threshold",
+        help="Override the score threshold for the 1st NMS operation",
+        type=float,
+    )
+    parser.add_argument(
+        "-t2",
+        "--second_nms_threshold",
+        help="Override the score threshold for the 2nd NMS operation",
+        type=float,
+    )
+    args = parser.parse_args()
+    if not all(
+        [
+            args.exported_onnx,
+            args.onnx,
+            args.det2_config,
+            args.det2_weights,
+            args.sample_image,
+        ]
+    ):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --exported_onnx --onnx --det2_config --det2_weights and --sample_image"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/eval_coco.py

@@ -0,0 +1,161 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 argparse
+import numpy as np
+import torch
+from PIL import Image
+from infer import TensorRTInfer
+from image_batcher import ImageBatcher
+
+try:
+    from detectron2.config import get_cfg
+    from detectron2.data import MetadataCatalog
+    from detectron2.evaluation import COCOEvaluator
+    from detectron2.structures import Instances, Boxes, ROIMasks
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+
+def build_evaluator(dataset_name):
+    """
+    Create evaluator for a COCO dataset.
+    Currently only Mask R-CNN is supported, dataset of interest is COCO, so only COCOEvaluator is implemented.
+    """
+    evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+    if evaluator_type in ["coco"]:
+        return COCOEvaluator(dataset_name)
+    else:
+        raise NotImplementedError("Evaluator type is not supported")
+
+
+def setup(config_file, weights):
+    """
+    Create config and perform basic setup.
+    """
+    cfg = get_cfg()
+    cfg.merge_from_file(config_file)
+    cfg.merge_from_list(["MODEL.WEIGHTS", weights])
+    cfg.freeze()
+    return cfg
+
+
+def main(args):
+    # Set up Detectron 2 config and build evaluator.
+    cfg = setup(args.det2_config, args.det2_weights)
+    dataset_name = cfg.DATASETS.TEST[0]
+    evaluator = build_evaluator(dataset_name)
+    evaluator.reset()
+
+    trt_infer = TensorRTInfer(args.engine)
+    batcher = ImageBatcher(
+        args.input, *trt_infer.input_spec(), config_file=args.det2_config
+    )
+
+    for batch, images, scales in batcher.get_batch():
+        print(
+            "Processing Image {} / {}".format(batcher.image_index, batcher.num_images),
+            end="\r",
+        )
+        detections = trt_infer.infer(batch, scales, args.nms_threshold)
+        for i in range(len(images)):
+            # Get inference image resolution.
+            infer_im = Image.open(images[i])
+            im_width, im_height = infer_im.size
+            pred_boxes = []
+            scores = []
+            pred_classes = []
+            # Number of detections.
+            num_instances = len(detections[i])
+            # Reserve numpy array to hold all mask predictions per image.
+            pred_masks = np.empty((num_instances, 28, 28), dtype=np.float32)
+            # Image ID, required for Detectron 2 evaluations.
+            source_id = int(os.path.splitext(os.path.basename(images[i]))[0])
+            # Loop over every single detection.
+            for n in range(num_instances):
+                det = detections[i][n]
+                # Append box coordinates data.
+                pred_boxes.append([det["ymin"], det["xmin"], det["ymax"], det["xmax"]])
+                # Append score.
+                scores.append(det["score"])
+                # Append class.
+                pred_classes.append(det["class"])
+                # Append mask.
+                pred_masks[n] = det["mask"]
+            # Create new Instances object required for Detectron 2 evalutions and add:
+            # boxes, scores, pred_classes, pred_masks.
+            image_shape = (im_height, im_width)
+            instances = Instances(image_shape)
+            instances.pred_boxes = Boxes(pred_boxes)
+            instances.scores = torch.tensor(scores)
+            instances.pred_classes = torch.tensor(pred_classes)
+            roi_masks = ROIMasks(torch.tensor(pred_masks))
+            instances.pred_masks = roi_masks.to_bitmasks(
+                instances.pred_boxes, im_height, im_width, args.iou_threshold
+            ).tensor
+            # Process evaluations per image.
+            image_dict = [{"instances": instances}]
+            input_dict = [{"image_id": source_id}]
+            evaluator.process(input_dict, image_dict)
+
+    # Final evaluations, generation of mAP accuracy performance.
+    evaluator.evaluate()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with.")
+    parser.add_argument(
+        "-i",
+        "--input",
+        help="The input to infer, either a single image path, or a directory of images.",
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-w", "--det2_weights", help="The Detectron 2 model weights (.pkl)", type=str
+    )
+    parser.add_argument(
+        "-t",
+        "--nms_threshold",
+        type=float,
+        help="Override the score threshold for the NMS operation, if higher than the threshold in the engine.",
+    )
+    parser.add_argument(
+        "--iou_threshold",
+        default=0.5,
+        type=float,
+        help="Select the IoU threshold for the mask segmentation. Range is 0 to 1. Pixel values more than threshold will become 1, less 0.",
+    )
+    args = parser.parse_args()
+    if not all([args.engine, args.input, args.det2_config, args.det2_weights]):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --engine --input --det2_config and --det2_weights"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/image_batcher.py

@@ -0,0 +1,222 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 numpy as np
+from PIL import Image
+
+try:
+    from detectron2.config import get_cfg
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+
+class ImageBatcher:
+    """
+    Creates batches of pre-processed images.
+    """
+
+    def __init__(
+        self,
+        input,
+        shape,
+        dtype,
+        max_num_images=None,
+        exact_batches=False,
+        config_file=None,
+    ):
+        """
+        :param input: The input directory to read images from.
+        :param shape: The tensor shape of the batch to prepare, either in NCHW or NHWC format.
+        :param dtype: The (numpy) datatype to cast the batched data to.
+        :param max_num_images: The maximum number of images to read from the directory.
+        :param exact_batches: This defines how to handle a number of images that is not an exact multiple of the batch
+        size. If false, it will pad the final batch with zeros to reach the batch size. If true, it will *remove* the
+        last few images in excess of a batch size multiple, to guarantee batches are exact (useful for calibration).
+        :param config_file: The path pointing to the Detectron 2 yaml file which describes the model.
+        """
+
+        def det2_setup(config_file):
+            """
+            Create configs and perform basic setups.
+            """
+            cfg = get_cfg()
+            if config_file is not None:
+                cfg.merge_from_file(config_file)
+            cfg.freeze()
+            return cfg
+
+        # Set up Detectron 2 model configuration.
+        self.det2_cfg = det2_setup(config_file)
+
+        # Extract min and max dimensions for testing.
+        self.min_size_test = self.det2_cfg.INPUT.MIN_SIZE_TEST
+        self.max_size_test = self.det2_cfg.INPUT.MAX_SIZE_TEST
+
+        # Find images in the given input path.
+        input = os.path.realpath(input)
+        self.images = []
+
+        extensions = [".jpg", ".jpeg", ".png", ".bmp", ".ppm"]
+
+        def is_image(path):
+            return (
+                os.path.isfile(path) and os.path.splitext(path)[1].lower() in extensions
+            )
+
+        if os.path.isdir(input):
+            self.images = [
+                os.path.join(input, f)
+                for f in os.listdir(input)
+                if is_image(os.path.join(input, f))
+            ]
+            self.images.sort()
+        elif os.path.isfile(input):
+            if is_image(input):
+                self.images.append(input)
+        self.num_images = len(self.images)
+        if self.num_images < 1:
+            print("No valid {} images found in {}".format("/".join(extensions), input))
+            sys.exit(1)
+
+        # Handle Tensor Shape.
+        self.dtype = dtype
+        self.shape = shape
+        assert len(self.shape) == 4
+        self.batch_size = shape[0]
+        assert self.batch_size > 0
+        self.format = None
+        self.width = -1
+        self.height = -1
+        if self.shape[1] == 3:
+            self.format = "NCHW"
+            self.height = self.shape[2]
+            self.width = self.shape[3]
+        elif self.shape[3] == 3:
+            self.format = "NHWC"
+            self.height = self.shape[1]
+            self.width = self.shape[2]
+        assert all([self.format, self.width > 0, self.height > 0])
+
+        # Adapt the number of images as needed.
+        if max_num_images and 0 < max_num_images < len(self.images):
+            self.num_images = max_num_images
+        if exact_batches:
+            self.num_images = self.batch_size * (self.num_images // self.batch_size)
+        if self.num_images < 1:
+            print("Not enough images to create batches")
+            sys.exit(1)
+        self.images = self.images[0 : self.num_images]
+
+        # Subdivide the list of images into batches.
+        self.num_batches = 1 + int((self.num_images - 1) / self.batch_size)
+        self.batches = []
+        for i in range(self.num_batches):
+            start = i * self.batch_size
+            end = min(start + self.batch_size, self.num_images)
+            self.batches.append(self.images[start:end])
+
+        # Indices.
+        self.image_index = 0
+        self.batch_index = 0
+
+    def preprocess_image(self, image_path):
+        """
+        The image preprocessor loads an image from disk and prepares it as needed for batching. This includes padding,
+        resizing, normalization, data type casting, and transposing.
+        This Image Batcher implements one algorithm for now:
+        * Resizes and pads the image to fit the input size.
+        :param image_path: The path to the image on disk to load.
+        :return: Two values: A numpy array holding the image sample, ready to be contacatenated into the rest of the
+        batch, and the resize scale used, if any.
+        """
+
+        def resize_pad(image, pad_color=(0, 0, 0)):
+            """
+            A subroutine to implement padding and resizing. This will resize the image to fit fully within the input
+            size, and pads the remaining bottom-right portions with the value provided.
+            :param image: The PIL image object
+            :pad_color: The RGB values to use for the padded area. Default: Black/Zeros.
+            :return: Two values: The PIL image object already padded and cropped, and the resize scale used.
+            """
+
+            # Get characteristics.
+            width, height = image.size
+
+            # Replicates behavior of ResizeShortestEdge augmentation.
+            size = self.min_size_test * 1.0
+            pre_scale = size / min(height, width)
+            if height < width:
+                newh, neww = size, pre_scale * width
+            else:
+                newh, neww = pre_scale * height, size
+
+            # If delta between min and max dimensions is so that max sized dimension reaches self.max_size_test
+            # before min dimension reaches self.min_size_test, keeping the same aspect ratio. We still need to
+            # maintain the same aspect ratio and keep max dimension at self.max_size_test.
+            if max(newh, neww) > self.max_size_test:
+                pre_scale = self.max_size_test * 1.0 / max(newh, neww)
+                newh = newh * pre_scale
+                neww = neww * pre_scale
+            neww = int(neww + 0.5)
+            newh = int(newh + 0.5)
+
+            # Scaling factor for normalized box coordinates scaling in post-processing.
+            scaling = max(newh / height, neww / width)
+
+            # Padding.
+            image = image.resize((neww, newh), resample=Image.BILINEAR)
+            pad = Image.new("RGB", (self.width, self.height))
+            pad.paste(pad_color, [0, 0, self.width, self.height])
+            pad.paste(image)
+            return pad, scaling
+
+        scale = None
+        image = Image.open(image_path)
+        image = image.convert(mode="RGB")
+        # Pad with mean values of COCO dataset, since padding is applied before actual model's
+        # preprocessor steps (Sub, Div ops), we need to pad with mean values in order to reverse
+        # the effects of Sub and Div, so that padding after model's preprocessor will be with actual 0s.
+        image, scale = resize_pad(image, (124, 116, 104))
+        image = np.asarray(image, dtype=np.float32)
+        # Change HWC -> CHW.
+        image = np.transpose(image, (2, 0, 1))
+        # Change RGB -> BGR.
+        return image[[2, 1, 0]], scale
+
+    def get_batch(self):
+        """
+        Retrieve the batches. This is a generator object, so you can use it within a loop as:
+        for batch, images in batcher.get_batch():
+           ...
+        Or outside of a batch with the next() function.
+        :return: A generator yielding three items per iteration: a numpy array holding a batch of images, the list of
+        paths to the images loaded within this batch, and the list of resize scales for each image in the batch.
+        """
+        for i, batch_images in enumerate(self.batches):
+            batch_data = np.zeros(self.shape, dtype=self.dtype)
+            batch_scales = [None] * len(batch_images)
+            for i, image in enumerate(batch_images):
+                self.image_index += 1
+                batch_data[i], batch_scales[i] = self.preprocess_image(image)
+            self.batch_index += 1
+            yield batch_data, batch_images, batch_scales

samples/python/detectron2/infer.py

@@ -0,0 +1,325 @@
+#
+# 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 argparse
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+from image_batcher import ImageBatcher
+from visualize import visualize_detections
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+
+class TensorRTInfer:
+    """
+    Implements inference for the Model TensorRT engine.
+    """
+
+    def __init__(self, engine_path):
+        """
+        :param engine_path: The path to the serialized engine to load from disk.
+        """
+
+        # Load TRT engine
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        trt.init_libnvinfer_plugins(self.logger, namespace="")
+        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
+            assert runtime
+            self.engine = runtime.deserialize_cuda_engine(f.read())
+        assert self.engine
+        self.context = self.engine.create_execution_context()
+        assert self.context
+
+        # Setup I/O bindings
+        self.inputs = []
+        self.outputs = []
+        self.device_memories = []
+        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 = self.engine.get_tensor_dtype(name)
+            shape = self.engine.get_tensor_shape(name)
+            if is_input:
+                self.batch_size = shape[0]
+            size = np.dtype(trt.nptype(dtype)).itemsize
+            for s in shape:
+                size *= s
+            device_mem = common.DeviceMem(size)
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": np.dtype(trt.nptype(dtype)),
+                "shape": list(shape),
+                "allocation": device_mem.device_ptr,
+                "size": size,
+            }
+            self.device_memories.append(device_mem)
+            if is_input:
+                self.inputs.append(binding)
+            else:
+                self.outputs.append(binding)
+
+        assert self.batch_size > 0
+        assert len(self.inputs) > 0
+        assert len(self.outputs) > 0
+        assert len(self.device_memories) > 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 infer(self, batch, scales=None, nms_threshold=None):
+        """
+        Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by
+        the ImageBatcher class. Memory copying to and from the GPU device will be performed here.
+        :param batch: A numpy array holding the image batch.
+        :param scales: The image resize scales for each image in this batch. Default: No scale postprocessing applied.
+        :return: A nested list for each image in the batch and each detection in the list.
+        """
+
+        # Prepare the output data.
+        outputs = []
+        for shape, dtype in self.output_spec():
+            outputs.append(np.zeros(shape, dtype))
+
+        # Process I/O and execute the network.
+        common.memcpy_host_to_device(
+            self.inputs[0]["allocation"], np.ascontiguousarray(batch)
+        )
+
+        self.context.execute_v2(self.allocations)
+        for o in range(len(outputs)):
+            common.memcpy_device_to_host(outputs[o], self.outputs[o]["allocation"])
+
+        # Process the results.
+        nums = outputs[0]
+        boxes = outputs[1]
+        scores = outputs[2]
+        pred_classes = outputs[3]
+        masks = outputs[4]
+
+        detections = []
+        for i in range(self.batch_size):
+            detections.append([])
+            for n in range(int(nums[i])):
+                # Select a mask.
+                mask = masks[i][n]
+
+                # Calculate scaling values for bboxes.
+                scale = self.inputs[0]["shape"][2]
+                scale /= scales[i]
+                scale_y = scale
+                scale_x = scale
+
+                if nms_threshold and scores[i][n] < nms_threshold:
+                    continue
+                # Append to detections
+                detections[i].append(
+                    {
+                        "ymin": boxes[i][n][0] * scale_y,
+                        "xmin": boxes[i][n][1] * scale_x,
+                        "ymax": boxes[i][n][2] * scale_y,
+                        "xmax": boxes[i][n][3] * scale_x,
+                        "score": scores[i][n],
+                        "class": int(pred_classes[i][n]),
+                        "mask": mask,
+                    }
+                )
+        return detections
+
+
+def main(args):
+    output_dir = os.path.realpath(args.output)
+    os.makedirs(output_dir, exist_ok=True)
+
+    labels = [
+        "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",
+    ]
+
+    trt_infer = TensorRTInfer(args.engine)
+    batcher = ImageBatcher(
+        args.input, *trt_infer.input_spec(), config_file=args.det2_config
+    )
+    for batch, images, scales in batcher.get_batch():
+        print(
+            "Processing Image {} / {}".format(batcher.image_index, batcher.num_images),
+            end="\r",
+        )
+        detections = trt_infer.infer(batch, scales, args.nms_threshold)
+        for i in range(len(images)):
+            basename = os.path.splitext(os.path.basename(images[i]))[0]
+            # Image Visualizations
+            output_path = os.path.join(output_dir, "{}.png".format(basename))
+            visualize_detections(
+                images[i], output_path, detections[i], labels, args.iou_threshold
+            )
+            # Text Results
+            output_results = ""
+            for d in detections[i]:
+                line = [
+                    d["xmin"],
+                    d["ymin"],
+                    d["xmax"],
+                    d["ymax"],
+                    d["score"],
+                    d["class"],
+                ]
+                output_results += "\t".join([str(f) for f in line]) + "\n"
+            with open(os.path.join(args.output, "{}.txt".format(basename)), "w") as f:
+                f.write(output_results)
+    print()
+    print("Finished Processing")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-e", "--engine", default=None, help="The serialized TensorRT engine"
+    )
+    parser.add_argument(
+        "-i", "--input", default=None, help="Path to the image or directory to process"
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="Directory where to save the visualization results",
+    )
+    parser.add_argument(
+        "-t",
+        "--nms_threshold",
+        type=float,
+        help="Override the score threshold for the NMS operation, if higher than the threshold in the engine.",
+    )
+    parser.add_argument(
+        "--iou_threshold",
+        default=0.5,
+        type=float,
+        help="Select the IoU threshold for the mask segmentation. Range is 0 to 1. Pixel values more than threshold will become 1, less 0",
+    )
+    args = parser.parse_args()
+    if not all([args.engine, args.input, args.output, args.det2_config]):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --engine --input --output and --det2_config"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/onnx_utils.py

@@ -0,0 +1,332 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 logging
+import numpy as np
+import onnx_graphsurgeon as gs
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("ModelHelper").setLevel(logging.INFO)
+log = logging.getLogger("ModelHelper")
+
+
+@gs.Graph.register()
+def op_with_const(self, op, name, input, value):
+    """
+    Add an operation with constant to the graph which will operate on the input tensor with the value(s) given.
+    :param op: The ONNX operation to perform, i.e. "Add" or "Mul".
+    :param input: The tensor to operate on.
+    :param value: The value array to operate with.
+    :param name: The name to use for the node.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}': {}".format(op, name, value.squeeze()))
+    const = gs.Constant(name="{}_value:0".format(name), values=value)
+    return self.layer(
+        name=name, op=op, inputs=[input_tensor, const], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def matmul(self, name, input, value):
+    """
+    Add MatMul operation to the graph which will operate on the input tensor with the value(s) given.
+    :param input: The tensor to operate on.
+    :param value: The linear transformation matrix to operate with.
+    :param name: The name to use for the node.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}': {}".format("MatMul", name, value.squeeze()))
+    const = gs.Constant(name="{}_value:0".format(name), values=value)
+    return self.layer(
+        name=name, op="MatMul", inputs=[input_tensor, const], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def clip(self, name, input, clip_min, clip_max):
+    """
+    Add Clip operation to the graph which will operate on the input tensor with the value(s) given.
+    :param input: The tensor to operate on.
+    :param name: The name to use for the node.
+    :param clip_min: Minimum value to include, less is clipped.
+    :param clip_max: Maximum value to include, more is clipped.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}".format("Clip", name))
+    const_min = gs.Constant(
+        name="{}_value:0".format(name), values=np.asarray([clip_min], dtype=np.float32)
+    )
+    const_max = gs.Constant(
+        name="{}_value:1".format(name), values=np.asarray([clip_max], dtype=np.float32)
+    )
+    return self.layer(
+        name=name,
+        op="Clip",
+        inputs=[input_tensor, const_min, const_max],
+        outputs=[name + ":0"],
+    )
+
+
+@gs.Graph.register()
+def slice(self, name, input, starts, ends, axes):
+    """
+    Add Slice operation to the graph which will operate on the input tensor with the value(s) given.
+    :param op: The ONNX operation to perform, i.e. "Add" or "Mul".
+    :param input: The tensor to operate on.
+    :param name: The name to use for the node.
+    :param starts: Value at which Slice starts.
+    :param ends: Value at which Slice ends.
+    :param axes: Axes on which Slice operation should be performed.
+    """
+
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}".format("Slice", name))
+    const_start = gs.Constant(
+        name="{}_value:0".format(name), values=np.asarray([starts], dtype=np.int64)
+    )
+    const_end = gs.Constant(
+        name="{}_value:1".format(name), values=np.asarray([ends], dtype=np.int64)
+    )
+    const_axes = gs.Constant(
+        name="{}_value:2".format(name), values=np.asarray([axes], dtype=np.int64)
+    )
+    return self.layer(
+        name=name,
+        op="Slice",
+        inputs=[input_tensor, const_start, const_end, const_axes],
+        outputs=[name + ":0"],
+    )
+
+
+@gs.Graph.register()
+def unsqueeze(self, name, input, axes=[3]):
+    """
+    Adds to the graph an Unsqueeze node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be "unsqueezed".
+    :param axes: A list of axes on which to add the new dimension(s).
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Unsqueeze node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Unsqueeze",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def squeeze(self, name, input, axes=[2]):
+    """
+    Adds to the graph an Squeeze node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be "squeezed".
+    :param axes: A list of axes on which to remove a dimension(s).
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Squeeze node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Squeeze",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def gather(self, name, data, indices, axes=0):
+    """
+    Adds to the graph a Gather node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param data: Data from which to gather specific tensors.
+    :param indices: Indices by which to gather data tensors.
+    :param axes: A list of axes on which to perform gather operation
+    """
+    data_tensor = data if type(data) is gs.Variable else data[0]
+    indices_tensor = indices if type(indices) is gs.Variable else indices[0]
+    log.debug("Created Gather node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Gather",
+        inputs=[data_tensor, indices_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def transpose(self, name, input, perm):
+    """
+    Adds to the graph a Transpose node for the given axes permutation and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be transposed.
+    :param perm: A list of axes defining their order after transposing occurs.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Transpose node '{}': {}".format(name, perm))
+    return self.layer(
+        name=name,
+        op="Transpose",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"perm": perm},
+    )
+
+
+@gs.Graph.register()
+def sigmoid(self, name, input):
+    """
+    Adds to the graph a Sigmoid node for the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be applied to.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Sigmoid node '{}'".format(name))
+    return self.layer(
+        name=name, op="Sigmoid", inputs=[input_tensor], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def plugin(self, op, name, inputs: list, outputs: list, attrs):
+    """
+    Adds to the graph a TensorRT plugin node with the given name, inputs and outputs. The attrs dictionary holds
+    attributes to be added to the plugin node.
+    :param self: The gs.Graph object being extended.
+    :param op: The registered name for the TensorRT plugin.
+    :param name: The name to use for the node.
+    :param inputs: The list of tensors to use an inputs.
+    :param outputs: The list of tensors to use as outputs.
+    :param attrs: The dictionary to use as attributes.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    log.debug("Created TRT Plugin node '{}': {}".format(name, attrs))
+    return self.layer(op=op, name=name, inputs=inputs, outputs=outputs, attrs=attrs)
+
+
+@gs.Graph.register()
+def find_node_by_op(self, op):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if node.op == op:
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_node_by_op_name(self, op, name):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :param name: Selected node name.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if node.op == op and node.name == name:
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_node_by_op_input_output_name(
+    self, op, input_name, output_name, input_pos=0, output_pos=0
+):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :param input_pos: Which input to consider, default is 0.
+    :param output_pos: Which output to consider, default is 0.
+    :param input_name: Selected input's name.
+    :param output_name: Selected output's name.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if (
+            node.op == op
+            and node.inputs[input_pos].name == input_name
+            and node.outputs[output_pos].name == output_name
+        ):
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_descendant_by_op(self, node, op, depth=10):
+    """
+    Starting from the given node, finds a node lower in the graph matching the given operation name.
+    This is not an exhaustive graph search.
+    In order to graph search bfs is used, so runtime complexity is O(V+E).
+    :param self: The gs.Graph object being extended.
+    :param node: The node to start searching from.
+    :param op: The operation name to search for.
+    :param depth: Stop searching after traversing these many nodes.
+    :return: The first descendant node matching that performs that op.
+    """
+    queue = []
+    for i in range(depth):
+        queue.append(node.o())
+        while queue:
+            node = queue.pop(0)
+            if node.op == op:
+                return node
+            for child in node.outputs[0].outputs:
+                queue.append(child)
+    return None
+
+
+@gs.Graph.register()
+def find_ancestor_by_op(self, node, op, depth=10):
+    """
+    Starting from the given node, finds a node higher in the graph matching the given operation name.
+    This is not an exhaustive graph search.
+    In order to graph search bfs is used, so runtime complexity is O(V+E).
+    :param self: The gs.Graph object being extended.
+    :param node: The node to start searching from.
+    :param op: The operation name to search for.
+    :param depth: Stop searching after traversing these many nodes.
+    :return: The first ancestor node matching that performs that op.
+    """
+    queue = []
+    for i in range(depth):
+        queue.append(node.i())
+        while queue:
+            node = queue.pop(0)
+            if node.op == op:
+                return node
+            for child in node.inputs[-1].inputs:
+                queue.append(child)
+    return None

samples/python/detectron2/requirements.txt

@@ -0,0 +1,12 @@
+onnx==1.18.0
+onnxruntime==1.18.1
+onnxscript>=0.6.0
+Pillow==11.3.0
+git+https://github.com/facebookresearch/detectron2.git
+git+https://github.com/NVIDIA/TensorRT#subdirectory=tools/onnx-graphsurgeon
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/detectron2/visualize.py

@@ -0,0 +1,255 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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
+import PIL.ImageFilter as ImageFilter
+
+
+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",
+]
+
+
+# Overlay mask with transparency on top of the image.
+def overlay(image, mask, color, alpha_transparency=0.5):
+    for channel in range(3):
+        image[:, :, channel] = np.where(
+            mask == 1,
+            image[:, :, channel] * (1 - alpha_transparency)
+            + alpha_transparency * color[channel] * 255,
+            image[:, :, channel],
+        )
+    return image
+
+
+def visualize_detections(
+    image_path, output_path, detections, labels=[], iou_threshold=0.5
+):
+    image = Image.open(image_path).convert(mode="RGB")
+    # Get image dimensions.
+    im_width, im_height = image.size
+    line_width = 2
+    font = ImageFont.load_default()
+    for d in detections:
+        color = COLORS[d["class"] % len(COLORS)]
+        # Dynamically convert PIL color into RGB numpy array.
+        pixel_color = Image.new("RGB", (1, 1), color)
+        # Normalize.
+        np_color = (np.asarray(pixel_color)[0][0]) / 255
+        # TRT instance segmentation masks.
+        if isinstance(d["mask"], np.ndarray) and d["mask"].shape == (28, 28):
+            # PyTorch uses [x1,y1,x2,y2] format instead of regular [y1,x1,y2,x2].
+            d["ymin"], d["xmin"], d["ymax"], d["xmax"] = (
+                d["xmin"],
+                d["ymin"],
+                d["xmax"],
+                d["ymax"],
+            )
+            # Get detection bbox resolution.
+            det_width = round(d["xmax"] - d["xmin"])
+            det_height = round(d["ymax"] - d["ymin"])
+            # Slight scaling, to get binary masks after float32 -> uint8
+            # conversion, if not scaled all pixels are zero.
+            mask = d["mask"] > iou_threshold
+            # Convert float32 -> uint8.
+            mask = mask.astype(np.uint8)
+            # Create an image out of predicted mask array.
+            small_mask = Image.fromarray(mask)
+            # Upsample mask to detection bbox's size.
+            mask = small_mask.resize((det_width, det_height), resample=Image.BILINEAR)
+            # Create an original image sized template for correct mask placement.
+            pad = Image.new("L", (im_width, im_height))
+            # Place your mask according to detection bbox placement.
+            pad.paste(mask, (round(d["xmin"]), (round(d["ymin"]))))
+            # Reconvert mask into numpy array for evaluation.
+            padded_mask = np.array(pad)
+            # Creat np.array from original image, copy in order to modify.
+            image_copy = np.asarray(image).copy()
+            # Image with overlaid mask.
+            masked_image = overlay(image_copy, padded_mask, np_color)
+            # Reconvert back to PIL.
+            image = Image.fromarray(masked_image)
+
+        # Bbox lines.
+        draw = ImageDraw.Draw(image)
+        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 = "Class {}".format(d["class"])
+        if d["class"] < len(labels):
+            label = "{}".format(labels[d["class"]])
+        score = d["score"]
+        text = "{}: {}%".format(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)