diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 6fa7e3b34..7cebac986 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -89,7 +89,7 @@ def inference_segmentor(model, img): # scatter to specified GPU data = scatter(data, [device])[0] else: - data['img_metas'] = data['img_metas'][0].data + data['img_metas'] = [i.data[0] for i in data['img_metas']] # forward the model with torch.no_grad(): diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 000000000..046e03628 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,29 @@ +import os.path as osp + +import mmcv + +from mmseg.apis import inference_segmentor, init_segmentor + + +def test_test_time_augmentation_on_cpu(): + config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py' + config = mmcv.Config.fromfile(config_file) + + # Remove pretrain model download for testing + config.model.pretrained = None + # Replace SyncBN with BN to inference on CPU + norm_cfg = dict(type='BN', requires_grad=True) + config.model.backbone.norm_cfg = norm_cfg + config.model.decode_head.norm_cfg = norm_cfg + config.model.auxiliary_head.norm_cfg = norm_cfg + + # Enable test time augmentation + config.data.test.pipeline[1].flip = True + + checkpoint_file = None + model = init_segmentor(config, checkpoint_file, device='cpu') + + img = mmcv.imread( + osp.join(osp.dirname(__file__), 'data/color.jpg'), 'color') + result = inference_segmentor(model, img) + assert result[0].shape == (288, 512)