Files
nvidia--tensorrt/plugin/common/mrcnn_config.h
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

178 lines
4.9 KiB
C++

/*
* 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.
*/
#ifndef TRT_PLUGIN_MASKRCNN_CONFIG_H
#define TRT_PLUGIN_MASKRCNN_CONFIG_H
#include "NvInfer.h"
#include <string>
#include <vector>
namespace MaskRCNNConfig
{
static const nvinfer1::Dims3 IMAGE_SHAPE{3, 1024, 1024};
// Pooled ROIs
static int32_t const POOL_SIZE = 7;
static int32_t const MASK_POOL_SIZE = 14;
// Threshold to determine the mask area out of final convolution output
static float const MASK_THRESHOLD = 0.5F;
// Bounding box refinement standard deviation for RPN and final detections.
static float const RPN_BBOX_STD_DEV[] = {0.1F, 0.1F, 0.2F, 0.2F};
static float const BBOX_STD_DEV[] = {0.1F, 0.1F, 0.2F, 0.2F};
// Max number of final detections
static int32_t const DETECTION_MAX_INSTANCES = 100;
// Minimum probability value to accept a detected instance
// ROIs below this threshold are skipped
static float const DETECTION_MIN_CONFIDENCE = 0.7F;
// Non-maximum suppression threshold for detection
static float const DETECTION_NMS_THRESHOLD = 0.3F;
// The strides of each layer of the FPN Pyramid. These values
// are based on a Resnet101 backbone.
static const std::vector<float> BACKBONE_STRIDES = {4.F, 8.F, 16.F, 32.F, 64.F};
// Size of the fully-connected layers in the classification graph
static int32_t const FPN_CLASSIF_FC_LAYERS_SIZE = 1024;
// Size of the top-down layers used to build the feature pyramid
static int32_t const TOP_DOWN_PYRAMID_SIZE = 256;
// Number of classification classes (including background)
static int32_t const NUM_CLASSES = 1 + 80; // COCO has 80 classes
// Length of square anchor side in pixels
static const std::vector<float> RPN_ANCHOR_SCALES = {32.F, 64.F, 128.F, 256.F, 512.F};
// Ratios of anchors at each cell (width/height)
// A value of 1 represents a square anchor, and 0.5 is a wide anchor
static float const RPN_ANCHOR_RATIOS[] = {0.5F, 1.F, 2.F};
// Anchor stride
// If 1 then anchors are created for each cell in the backbone feature map.
// If 2, then anchors are created for every other cell, and so on.
static int32_t const RPN_ANCHOR_STRIDE = 1;
// Although Python impementation uses 6000,
// TRT fails if this number larger than kMAX_TOPK_K defined in engine/checkMacros.h
static int32_t const MAX_PRE_NMS_RESULTS = 1024; // 3840;
// Non-max suppression threshold to filter RPN proposals.
// You can increase this during training to generate more propsals.
static float const RPN_NMS_THRESHOLD = 0.7F;
// ROIs kept after non-maximum suppression (training and inference)
static int32_t const POST_NMS_ROIS_INFERENCE = 1000;
// COCO Class names
static const std::vector<std::string> CLASS_NAMES = {
"BG",
"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",
};
static const std::string MODEL_NAME = "mrcnn_nchw.uff";
static const std::string MODEL_INPUT = "input_image";
static const nvinfer1::Dims3 MODEL_INPUT_SHAPE = IMAGE_SHAPE;
static const std::vector<std::string> MODEL_OUTPUTS = {"mrcnn_detection", "mrcnn_mask/Sigmoid"};
static const nvinfer1::Dims2 MODEL_DETECTION_SHAPE{DETECTION_MAX_INSTANCES, 6};
static const nvinfer1::Dims4 MODEL_MASK_SHAPE{DETECTION_MAX_INSTANCES, NUM_CLASSES, 28, 28};
} // namespace MaskRCNNConfig
#endif // TRT_PLUGIN_MASKRCNN_CONFIG_H