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
(self, anchors, first_nms_threshold=None, second_nms_threshold=None)
| 343 | return roi_align_output |
| 344 | |
| 345 | def process_graph(self, anchors, first_nms_threshold=None, second_nms_threshold=None): |
| 346 | """ |
| 347 | Processes the graph to replace the GenerateProposals and BoxWithNMSLimit operations with EfficientNMS_TRT |
| 348 | TensorRT plugin nodes and ROIAlign operations with PyramidROIAlign_TRT plugin nodes. |
| 349 | :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided |
| 350 | inside the graph. |
| 351 | :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph. |
| 352 | :param second_nms_threshold: Override the 2nd NMS score threshold value. If set to None, use the value in the graph. |
| 353 | """ |
| 354 | def backbone(): |
| 355 | """ |
| 356 | Updates the graph to replace all ResizeNearest ops with ResizeNearest plugins in backbone. |
| 357 | """ |
| 358 | # Get final backbone outputs. |
| 359 | p2 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output2/Conv") |
| 360 | p3 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output3/Conv") |
| 361 | p4 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output4/Conv") |
| 362 | p5 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output5/Conv") |
| 363 | |
| 364 | |
| 365 | return p2.outputs[0], p3.outputs[0], p4.outputs[0], p5.outputs[0] |
| 366 | |
| 367 | def proposal_generator(anchors, first_nms_threshold): |
| 368 | """ |
| 369 | Updates the graph to replace all GenerateProposals Caffe ops with one single NMS for proposals generation. |
| 370 | :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided |
| 371 | inside the graph |
| 372 | :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph. |
| 373 | """ |
| 374 | # Get nodes containing final objectness logits. |
| 375 | p2_logits = self.graph.find_node_by_op_name("Flatten", "/proposal_generator/Flatten") |
| 376 | p3_logits = self.graph.find_node_by_op_name("Flatten", "/proposal_generator/Flatten_1") |
| 377 | p4_logits = self.graph.find_node_by_op_name("Flatten", "/proposal_generator/Flatten_2") |
| 378 | p5_logits = self.graph.find_node_by_op_name("Flatten", "/proposal_generator/Flatten_3") |
| 379 | p6_logits = self.graph.find_node_by_op_name("Flatten", "/proposal_generator/Flatten_4") |
| 380 | |
| 381 | # Get nodes containing final anchor_deltas. |
| 382 | p2_anchors = self.graph.find_node_by_op_name("Reshape", "/proposal_generator/Reshape_1") |
| 383 | p3_anchors = self.graph.find_node_by_op_name("Reshape", "/proposal_generator/Reshape_3") |
| 384 | p4_anchors = self.graph.find_node_by_op_name("Reshape", "/proposal_generator/Reshape_5") |
| 385 | p5_anchors = self.graph.find_node_by_op_name("Reshape", "/proposal_generator/Reshape_7") |
| 386 | p6_anchors = self.graph.find_node_by_op_name("Reshape", "/proposal_generator/Reshape_9") |
| 387 | |
| 388 | # Concatenate all objectness logits/scores data. |
| 389 | scores_inputs = [p2_logits.outputs[0], p3_logits.outputs[0], p4_logits.outputs[0], p5_logits.outputs[0], p6_logits.outputs[0]] |
| 390 | scores_tensor = self.graph.layer(name="scores", op="Concat", inputs=scores_inputs, outputs=['scores'], attrs={'axis': 1})[0] |
| 391 | # Unsqueeze to add 3rd dimension of 1 to match tensor dimensions of boxes tensor. |
| 392 | scores = self.graph.unsqueeze("scores_unsqueeze", scores_tensor, [2])[0] |
| 393 | |
| 394 | # Concatenate all boxes/anchor_delta data. |
| 395 | boxes_inputs = [p2_anchors.outputs[0], p3_anchors.outputs[0], p4_anchors.outputs[0], p5_anchors.outputs[0], p6_anchors.outputs[0]] |
| 396 | boxes = self.graph.layer(name="boxes", op="Concat", inputs=boxes_inputs, outputs=['anchors'], attrs={'axis': 1})[0] |
| 397 | |
| 398 | # Convert the anchors from Corners to CenterSize encoding. |
| 399 | anchors = np.matmul(anchors, [[0.5, 0, -1, 0], [0, 0.5, 0, -1], [0.5, 0, 1, 0], [0, 0.5, 0, 1]]) |
| 400 | anchors = anchors / [self.width, self.height, self.width, self.height] # Normalize anchors to [0-1] range |
| 401 | anchors = np.expand_dims(anchors, axis=0) |
| 402 | anchors = anchors.astype(np.float32) |