Method NMS

samples/python/detectron2/create_onnx.py:245–298

(self, boxes, scores, anchors, background_class, score_activation, max_proposals, iou_threshold, nms_score_threshold, user_threshold, nms_name=None)

                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.