| 298 | return nms_outputs |
| 299 | |
| 300 | def ROIAlign(self, rois, p2, p3, p4, p5, pooled_size, sampling_ratio, roi_align_type, num_rois, ra_name): |
| 301 | # Helper function to create the ROIAlign Plugin node with the selected inputs. |
| 302 | # PyramidROIAlign_TRT TensorRT Plugin is suitable for our use case. |
| 303 | # :param rois: Regions of interest/detection boxes outputs from preceding NMS node. |
| 304 | # :param p2: Output of p2 feature map. |
| 305 | # :param p3: Output of p3 feature map. |
| 306 | # :param p4: Output of p4 feature map. |
| 307 | # :param p5: Output of p5 feature map. |
| 308 | # :param pooled_size: Pooled output dimensions. |
| 309 | # :param sampling_ratio: Number of sampling points in the interpolation grid used to compute the output value of each pooled output bin. |
| 310 | # :param roi_align_type: Type of Detectron 2 ROIAlign op, either ROIAlign (vanilla) or ROIAlignV2 (0.5 coordinate offset). |
| 311 | # :param num_rois: Number of ROIs resulting from ROIAlign operation. |
| 312 | # :param ra_name: Name of ROIAlign node in a graph, renames ROIAlign elements accordingly in order to eliminate cycles. |
| 313 | |
| 314 | # Different types of Detectron 2's ROIAlign ops require coordinate offset that is supported by PyramidROIAlign_TRT. |
| 315 | if roi_align_type == "ROIAlignV2": |
| 316 | roi_coords_transform = 2 |
| 317 | elif roi_align_type == "ROIAlign": |
| 318 | roi_coords_transform = 0 |
| 319 | |
| 320 | # ROIAlign outputs. |
| 321 | roi_align_output = gs.Variable(name="roi_align/output_"+ra_name, dtype=np.float32, |
| 322 | shape=[self.batch_size, num_rois, self.fpn_out_channels, pooled_size, pooled_size]) |
| 323 | |
| 324 | # Plugin. |
| 325 | self.graph.plugin( |
| 326 | op="PyramidROIAlign_TRT", |
| 327 | name="roi_align_"+ra_name, |
| 328 | inputs=[rois, p2, p3, p4, p5], |
| 329 | outputs=[roi_align_output], |
| 330 | attrs={ |
| 331 | 'plugin_version': "1", |
| 332 | 'fpn_scale': 224, |
| 333 | 'pooled_size': pooled_size, |
| 334 | 'image_size': [self.height, self.width], |
| 335 | 'roi_coords_absolute': 0, |
| 336 | 'roi_coords_swap': 0, |
| 337 | 'roi_coords_transform': roi_coords_transform, |
| 338 | 'sampling_ratio': sampling_ratio, |
| 339 | } |
| 340 | ) |
| 341 | log.info("Created {} with PyramidROIAlign_TRT plugin".format(ra_name)) |
| 342 | |
| 343 | return roi_align_output |
| 344 | |
| 345 | def process_graph(self, anchors, first_nms_threshold=None, second_nms_threshold=None): |
| 346 | """ |