Detectron 2 exported ONNX does not contain anchors required for efficientNMS plug-in, so they must be generated "offline" by calling actual Detectron 2 model and getting anchors from it. :param sample_image: Sample image required to run through the model and obtain anchors.
(self, sample_image)
| 142 | break |
| 143 | |
| 144 | def get_anchors(self, sample_image): |
| 145 | """ |
| 146 | Detectron 2 exported ONNX does not contain anchors required for efficientNMS plug-in, so they must be generated |
| 147 | "offline" by calling actual Detectron 2 model and getting anchors from it. |
| 148 | :param sample_image: Sample image required to run through the model and obtain anchors. |
| 149 | Can be any image from a dataset. Make sure listed here Detectron 2 preprocessing steps |
| 150 | actually match your preprocessing steps. Otherwise, behavior can be unpredictable. |
| 151 | Additionally, anchors have to be generated for a fixed input dimensions, |
| 152 | meaning as soon as image leaves a preprocessor and enters predictor.model.backbone() it must have |
| 153 | a fixed dimension (1344x1344 in my case) that every single image in dataset must follow, since currently |
| 154 | TensorRT plug-ins do not support dynamic shapes. |
| 155 | """ |
| 156 | # Get Detectron 2 model config and build it. |
| 157 | predictor = DefaultPredictor(self.det2_cfg) |
| 158 | model = build_model(self.det2_cfg) |
| 159 | |
| 160 | # Image preprocessing. |
| 161 | input_im = cv2.imread(sample_image) |
| 162 | raw_height, raw_width = input_im.shape[:2] |
| 163 | image = predictor.aug.get_transform(input_im).apply_image(input_im) |
| 164 | image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) |
| 165 | |
| 166 | # Model preprocessing. |
| 167 | inputs = [{"image": image, "height": raw_height, "width": raw_width}] |
| 168 | images = [x["image"].to(model.device) for x in inputs] |
| 169 | images = [(x - model.pixel_mean) / model.pixel_std for x in images] |
| 170 | imagelist_images = ImageList.from_tensors(images, 1344) |
| 171 | |
| 172 | # Get feature maps from backbone. |
| 173 | features = predictor.model.backbone(imagelist_images.tensor) |
| 174 | |
| 175 | # Get proposals from Region Proposal Network and obtain anchors from anchor generator. |
| 176 | features = [features[f] for f in predictor.model.proposal_generator.in_features] |
| 177 | det2_anchors = predictor.model.proposal_generator.anchor_generator(features) |
| 178 | |
| 179 | # Extract anchors based on feature maps in ascending order (P2->P6). |
| 180 | p2_anchors = det2_anchors[0].tensor.detach().cpu().numpy() |
| 181 | p3_anchors = det2_anchors[1].tensor.detach().cpu().numpy() |
| 182 | p4_anchors = det2_anchors[2].tensor.detach().cpu().numpy() |
| 183 | p5_anchors = det2_anchors[3].tensor.detach().cpu().numpy() |
| 184 | p6_anchors = det2_anchors[4].tensor.detach().cpu().numpy() |
| 185 | final_anchors = np.concatenate((p2_anchors,p3_anchors,p4_anchors,p5_anchors,p6_anchors)) |
| 186 | |
| 187 | return final_anchors |
| 188 | |
| 189 | def save(self, output_path): |
| 190 | """ |