Simple test with single image.
(self, img, img_meta, rescale=True)
| 245 | return output |
| 246 | |
| 247 | def simple_test(self, img, img_meta, rescale=True): |
| 248 | """Simple test with single image.""" |
| 249 | seg_logit = self.inference(img, img_meta, rescale) |
| 250 | if self.out_channels == 1: |
| 251 | seg_pred = (seg_logit > self.decode_head.threshold).to(seg_logit).squeeze(1) |
| 252 | else: |
| 253 | seg_pred = seg_logit.argmax(dim=1) |
| 254 | if torch.onnx.is_in_onnx_export(): |
| 255 | # our inference backend only support 4D output |
| 256 | seg_pred = seg_pred.unsqueeze(0) |
| 257 | return seg_pred |
| 258 | seg_pred = seg_pred.cpu().numpy() |
| 259 | # unravel batch dim |
| 260 | seg_pred = list(seg_pred) |
| 261 | return seg_pred |
| 262 | |
| 263 | def simple_test_logits(self, img, img_metas, rescale=True): |
| 264 | """Test without augmentations. |