mmdetection/tests/test_models/test_roi_heads/test_standard_roi_head.py

# Copyright (c) OpenMMLab. All rights reserved.
import unittest
from unittest import TestCase

import torch
from mmengine.config import Config
from parameterized import parameterized

from mmdet.registry import MODELS
from mmdet.testing import demo_mm_inputs, demo_mm_proposals
from mmdet.utils import register_all_modules

register_all_modules()


def _fake_roi_head(with_shared_head=False):
    """Set a fake roi head config."""
    if not with_shared_head:
        roi_head = Config(
            dict(
                type='StandardRoIHead',
                bbox_roi_extractor=dict(
                    type='SingleRoIExtractor',
                    roi_layer=dict(
                        type='RoIAlign', output_size=7, sampling_ratio=0),
                    out_channels=1,
                    featmap_strides=[4, 8, 16, 32]),
                bbox_head=dict(
                    type='Shared2FCBBoxHead',
                    in_channels=1,
                    fc_out_channels=1,
                    num_classes=4),
                mask_roi_extractor=dict(
                    type='SingleRoIExtractor',
                    roi_layer=dict(
                        type='RoIAlign', output_size=14, sampling_ratio=0),
                    out_channels=1,
                    featmap_strides=[4, 8, 16, 32]),
                mask_head=dict(
                    type='FCNMaskHead',
                    num_convs=1,
                    in_channels=1,
                    conv_out_channels=1,
                    num_classes=4),
                train_cfg=dict(
                    assigner=dict(
                        type='MaxIoUAssigner',
                        pos_iou_thr=0.5,
                        neg_iou_thr=0.5,
                        min_pos_iou=0.5,
                        match_low_quality=True,
                        ignore_iof_thr=-1),
                    sampler=dict(
                        type='RandomSampler',
                        num=512,
                        pos_fraction=0.25,
                        neg_pos_ub=-1,
                        add_gt_as_proposals=True),
                    mask_size=28,
                    pos_weight=-1,
                    debug=False),
                test_cfg=dict(
                    score_thr=0.05,
                    nms=dict(type='nms', iou_threshold=0.5),
                    max_per_img=100,
                    mask_thr_binary=0.5)))
    else:
        roi_head = Config(
            dict(
                type='StandardRoIHead',
                shared_head=dict(
                    type='ResLayer',
                    depth=50,
                    stage=3,
                    stride=2,
                    dilation=1,
                    style='caffe',
                    norm_cfg=dict(type='BN', requires_grad=False),
                    norm_eval=True),
                bbox_roi_extractor=dict(
                    type='SingleRoIExtractor',
                    roi_layer=dict(
                        type='RoIAlign', output_size=14, sampling_ratio=0),
                    out_channels=1,
                    featmap_strides=[16]),
                bbox_head=dict(
                    type='BBoxHead',
                    with_avg_pool=True,
                    in_channels=2048,
                    roi_feat_size=7,
                    num_classes=4),
                mask_roi_extractor=None,
                mask_head=dict(
                    type='FCNMaskHead',
                    num_convs=0,
                    in_channels=2048,
                    conv_out_channels=1,
                    num_classes=4),
                train_cfg=dict(
                    assigner=dict(
                        type='MaxIoUAssigner',
                        pos_iou_thr=0.5,
                        neg_iou_thr=0.5,
                        min_pos_iou=0.5,
                        match_low_quality=False,
                        ignore_iof_thr=-1),
                    sampler=dict(
                        type='RandomSampler',
                        num=512,
                        pos_fraction=0.25,
                        neg_pos_ub=-1,
                        add_gt_as_proposals=True),
                    mask_size=14,
                    pos_weight=-1,
                    debug=False),
                test_cfg=dict(
                    score_thr=0.05,
                    nms=dict(type='nms', iou_threshold=0.5),
                    max_per_img=100,
                    mask_thr_binary=0.5)))
    return roi_head


class TestStandardRoIHead(TestCase):

    def test_init(self):
        """Test init standard RoI head."""
        # Normal Mask R-CNN RoI head
        roi_head_cfg = _fake_roi_head()
        roi_head = MODELS.build(roi_head_cfg)
        self.assertTrue(roi_head.with_bbox)
        self.assertTrue(roi_head.with_mask)

        # Mask R-CNN RoI head with shared_head
        roi_head_cfg = _fake_roi_head(with_shared_head=True)
        roi_head = MODELS.build(roi_head_cfg)
        self.assertTrue(roi_head.with_bbox)
        self.assertTrue(roi_head.with_mask)
        self.assertTrue(roi_head.with_shared_head)

    @parameterized.expand([(False, ), (True, )])
    def test_standard_roi_head_loss(self, with_shared_head):
        """Tests standard roi head loss when truth is empty and non-empty."""
        if not torch.cuda.is_available():
            # RoI pooling only support in GPU
            return unittest.skip('test requires GPU and torch+cuda')
        s = 256
        roi_head_cfg = _fake_roi_head(with_shared_head=with_shared_head)
        roi_head = MODELS.build(roi_head_cfg)
        roi_head = roi_head.cuda()
        feats = []
        for i in range(len(roi_head.bbox_roi_extractor.featmap_strides)):
            if not with_shared_head:
                feats.append(
                    torch.rand(1, 1, s // (2**(i + 2)),
                               s // (2**(i + 2))).to(device='cuda'))
            else:
                feats.append(
                    torch.rand(1, 1024, s // (2**(i + 2)),
                               s // (2**(i + 2))).to(device='cuda'))
        feats = tuple(feats)

        # When truth is non-empty then both cls, box, and mask loss
        # should be nonzero for random inputs
        image_shapes = [(3, s, s)]
        batch_data_samples = demo_mm_inputs(
            batch_size=1,
            image_shapes=image_shapes,
            num_items=[1],
            num_classes=4,
            with_mask=True,
            device='cuda')['data_samples']
        proposals_list = demo_mm_proposals(
            image_shapes=image_shapes, num_proposals=100, device='cuda')

        out = roi_head.loss(feats, proposals_list, batch_data_samples)
        loss_cls = out['loss_cls']
        loss_bbox = out['loss_bbox']
        loss_mask = out['loss_mask']
        self.assertGreater(loss_cls.sum(), 0, 'cls loss should be non-zero')
        self.assertGreater(loss_bbox.sum(), 0, 'box loss should be non-zero')
        self.assertGreater(loss_mask.sum(), 0, 'mask loss should be non-zero')

        # When there is no truth, the cls loss should be nonzero but
        # there should be no box and mask loss.
        batch_data_samples = demo_mm_inputs(
            batch_size=1,
            image_shapes=image_shapes,
            num_items=[0],
            num_classes=4,
            with_mask=True,
            device='cuda')['data_samples']
        proposals_list = demo_mm_proposals(
            image_shapes=image_shapes, num_proposals=100, device='cuda')
        out = roi_head.loss(feats, proposals_list, batch_data_samples)
        empty_cls_loss = out['loss_cls']
        empty_bbox_loss = out['loss_bbox']
        empty_mask_loss = out['loss_mask']
        self.assertGreater(empty_cls_loss.sum(), 0,
                           'cls loss should be non-zero')
        self.assertEqual(
            empty_bbox_loss.sum(), 0,
            'there should be no box loss when there are no true boxes')
        self.assertEqual(
            empty_mask_loss.sum(), 0,
            'there should be no mask loss when there are no true boxes')