| 45 | |
| 46 | |
| 47 | class DET2GraphSurgeon: |
| 48 | def __init__(self, saved_model_path, config_file, weights): |
| 49 | """ |
| 50 | Constructor of the Model Graph Surgeon object, to do the conversion of a Detectron 2 Mask R-CNN exported model |
| 51 | to an ONNX-TensorRT parsable model. |
| 52 | :param saved_model_path: The path pointing to the exported Detectron 2 Mask R-CNN ONNX model. |
| 53 | :param config_file: The path pointing to the Detectron 2 yaml file which describes the model. |
| 54 | :param config_file: Weights to load for the Detectron 2 model. |
| 55 | """ |
| 56 | |
| 57 | def det2_setup(config_file, weights): |
| 58 | """ |
| 59 | Create configs and perform basic setups. |
| 60 | """ |
| 61 | cfg = get_cfg() |
| 62 | cfg.merge_from_file(config_file) |
| 63 | cfg.merge_from_list(["MODEL.WEIGHTS", weights]) |
| 64 | cfg.freeze() |
| 65 | return cfg |
| 66 | |
| 67 | # Import exported Detectron 2 Mask R-CNN ONNX model as GraphSurgeon object. |
| 68 | self.graph = gs.import_onnx(onnx.load(saved_model_path)) |
| 69 | assert self.graph |
| 70 | log.info("ONNX graph loaded successfully") |
| 71 | |
| 72 | # Fold constants via ONNX-GS that exported script might've missed. |
| 73 | self.graph.fold_constants() |
| 74 | |
| 75 | # Set up Detectron 2 model configuration. |
| 76 | self.det2_cfg = det2_setup(config_file, weights) |
| 77 | |
| 78 | # Getting model characteristics. |
| 79 | self.fpn_out_channels = self.det2_cfg.MODEL.FPN.OUT_CHANNELS |
| 80 | self.num_classes = self.det2_cfg.MODEL.ROI_HEADS.NUM_CLASSES |
| 81 | self.first_NMS_max_proposals = self.det2_cfg.MODEL.RPN.POST_NMS_TOPK_TEST |
| 82 | self.first_NMS_iou_threshold = self.det2_cfg.MODEL.RPN.NMS_THRESH |
| 83 | self.first_NMS_score_threshold = 0.01 |
| 84 | self.first_ROIAlign_pooled_size = self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION |
| 85 | self.first_ROIAlign_sampling_ratio = self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO |
| 86 | self.first_ROIAlign_type = self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE |
| 87 | self.second_NMS_max_proposals = self.det2_cfg.TEST.DETECTIONS_PER_IMAGE |
| 88 | self.second_NMS_iou_threshold = self.det2_cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST |
| 89 | self.second_NMS_score_threshold = self.det2_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST |
| 90 | self.second_ROIAlign_pooled_size = self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION |
| 91 | self.second_ROIAlign_sampling_ratio = self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_SAMPLING_RATIO |
| 92 | self.second_ROIAlign_type = self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_TYPE |
| 93 | self.mask_out_res = 28 |
| 94 | |
| 95 | # Model characteristics. |
| 96 | log.info("Number of FPN output channels is {}".format(self.fpn_out_channels)) |
| 97 | log.info("Number of classes is {}".format(self.num_classes)) |
| 98 | log.info("First NMS max proposals is {}".format(self.first_NMS_max_proposals)) |
| 99 | log.info("First NMS iou threshold is {}".format(self.first_NMS_iou_threshold)) |
| 100 | log.info("First NMS score threshold is {}".format(self.first_NMS_score_threshold)) |
| 101 | log.info("First ROIAlign type is {}".format(self.first_ROIAlign_type)) |
| 102 | log.info("First ROIAlign pooled size is {}".format(self.first_ROIAlign_pooled_size)) |
| 103 | log.info("First ROIAlign sampling ratio is {}".format(self.first_ROIAlign_sampling_ratio)) |
| 104 | log.info("Second NMS max proposals is {}".format(self.second_NMS_max_proposals)) |