mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

add support for 4D output (#150)

Browse Source
This commit is contained in:
robin Han 2020-09-23 17:01:20 +08:00 committed by GitHub
parent a2738fd9be
commit 08f30ea497
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mmseg/models/segmentors/encoder_decoder.py
View File

@ -265,6 +265,8 @@ class EncoderDecoder(BaseSegmentor):
seg_logit = self.inference(img, img_meta, rescale) seg_logit = self.inference(img, img_meta, rescale)
seg_pred = seg_logit.argmax(dim=1) seg_pred = seg_logit.argmax(dim=1)
if torch.onnx.is_in_onnx_export(): if torch.onnx.is_in_onnx_export():
# our inference backend only support 4D output
seg_pred = seg_pred.unsqueeze(0)
return seg_pred return seg_pred
seg_pred = seg_pred.cpu().numpy() seg_pred = seg_pred.cpu().numpy()
# unravel batch dim # unravel batch dim