Method roi_heads

samples/python/detectron2/create_onnx.py:410–511

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

(rpn_outputs, p2, p3, p4, p5, second_nms_threshold)

            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)