Maskrcnn Error during inference when test time augmentation is on (KeyError 'labels')

🐛 Bug

Error during inference for segmentation models (MaskPostProcessor) when test time augmentation is turned on. This is because in the box_head/inference.py PostProcessor module during forward the post processor skips filtering the results if bbox_aug_enabled is True and so there is no "labels" field in the MaskPostProcessor input:

if not self.bbox_aug_enabled:  # If bbox aug is enabled, we will do it later
    boxlist = self.filter_results(boxlist, num_classes)

Also because of skipping filtering there is risk that you can run out of memory if there are too many proposals. My config file is as follows (modified e2e_mask_rcnn_R_50_FPN_1x.yaml to include TTA, 1 image per batch and smaller image size):

To Reproduce

Steps to reproduce the behavior:

Modify e2e_mask_rcnn_R_50_FPN_1x.yaml to include TTA, and depending on memory available allow 1 image per batch and smaller image size.
python test_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml"

MODEL:
  META_ARCHITECTURE: "GeneralizedRCNN"
  WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"
  BACKBONE:
    CONV_BODY: "R-50-FPN"
  RESNETS:
    BACKBONE_OUT_CHANNELS: 256
  RPN:
    USE_FPN: True
    ANCHOR_STRIDE: (4, 8, 16, 32, 64)
    PRE_NMS_TOP_N_TRAIN: 2000
    PRE_NMS_TOP_N_TEST: 1000
    POST_NMS_TOP_N_TEST: 1000
    FPN_POST_NMS_TOP_N_TEST: 1000
  ROI_HEADS:
    USE_FPN: True
  ROI_BOX_HEAD:
    POOLER_RESOLUTION: 7
    POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
    POOLER_SAMPLING_RATIO: 2
    FEATURE_EXTRACTOR: "FPN2MLPFeatureExtractor"
    PREDICTOR: "FPNPredictor"
  ROI_MASK_HEAD:
    POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
    FEATURE_EXTRACTOR: "MaskRCNNFPNFeatureExtractor"
    PREDICTOR: "MaskRCNNC4Predictor"
    POOLER_RESOLUTION: 14
    POOLER_SAMPLING_RATIO: 2
    RESOLUTION: 28
    SHARE_BOX_FEATURE_EXTRACTOR: False
  MASK_ON: True
DATASETS:
  TRAIN: ("coco_2014_train", "coco_2014_valminusminival")
  TEST: ("coco_2017_val",)
DATALOADER:
  SIZE_DIVISIBILITY: 32
  NUM_WORKERS: 0
SOLVER:
  BASE_LR: 0.02
  WEIGHT_DECAY: 0.0001
  STEPS: (60000, 80000)
  MAX_ITER: 90000
TEST:
   BBOX_AUG:
     ENABLED: True
     H_FLIP: True
     SCALE_H_FLIP: True
   IMS_PER_BATCH: 1
INPUT:
  MIN_SIZE_TEST: 10
  MAX_SIZE_TEST: 20

Traceback (most recent call last):
  File "test_net.py", line 116, in <module>
    main()
  File "test_net.py", line 108, in main
    output_folder=output_folder,
  File "maskrcnn-benchmark/maskrcnn_benchmark/engine/inference.py", line 84, in inference
    predictions = compute_on_dataset(model, data_loader, device, inference_timer)
  File "maskrcnn-benchmark/maskrcnn_benchmark/engine/inference.py", line 28, in compute_on_dataset
    output = im_detect_bbox_aug(model, images, device)
  File "maskrcnn-benchmark/maskrcnn_benchmark/engine/bbox_aug.py", line 28, in im_detect_bbox_aug
    model, images, cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MAX_SIZE_TEST, device
  File "maskrcnn-benchmark/maskrcnn_benchmark/engine/bbox_aug.py", line 84, in im_detect_bbox
    return model(images.to(device))
  File "miniconda2/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "maskrcnn-benchmark/maskrcnn_benchmark/modeling/detector/generalized_rcnn.py", line 52, in forward
    x, result, detector_losses = self.roi_heads(features, proposals, targets)
  File "miniconda2/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "maskrcnn-benchmark/maskrcnn_benchmark/modeling/roi_heads/roi_heads.py", line 39, in forward
    x, detections, loss_mask = self.mask(mask_features, detections, targets)
  File "miniconda2/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "maskrcnn-benchmark/maskrcnn_benchmark/modeling/roi_heads/mask_head/mask_head.py", line 74, in forward
    result = self.post_processor(mask_logits, proposals)
  File "miniconda2/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "maskrcnn-benchmark/maskrcnn_benchmark/modeling/roi_heads/mask_head/inference.py", line 42, in forward
    labels = [bbox.get_field("labels") for bbox in boxes]
  File "maskrcnn-benchmark/maskrcnn_benchmark/modeling/roi_heads/mask_head/inference.py", line 42, in <listcomp>
    labels = [bbox.get_field("labels") for bbox in boxes]
  File "maskrcnn-benchmark/maskrcnn_benchmark/structures/bounding_box.py", line 43, in get_field
    return self.extra_fields[field]
KeyError: 'labels'

Expected behavior

Maskrcnn model during inference and tta enabled to return normally.

Environment

Collecting environment information... PyTorch version: 1.1.0 Is debug build: No CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.2 LTS GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0 CMake version: version 3.10.2

Python version: 3.7 Is CUDA available: Yes CUDA runtime version: 10.0.130 GPU models and configuration: GPU 0: Quadro P5000 Nvidia driver version: 418.67 cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.2

Versions of relevant libraries: [pip] numpy==1.16.4 [pip] torch==1.1.0 [pip] torchvision==0.3.0 [conda] blas 1.0 mkl
[conda] mkl 2019.4 243
[conda] mkl_fft 1.0.12 py37ha843d7b_0
[conda] mkl_random 1.0.2 py37hd81dba3_0
[conda] pytorch 1.1.0 py3.7_cuda10.0.130_cudnn7.5.1_0 pytorch [conda] torchvision 0.3.0 py37_cu10.0.130_1 pytorch

Additional context

Also refer to issue here

facebookresearch / maskrcnn-benchmark