base_3droi_head.py 2.92 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
wuyuefeng's avatar
wuyuefeng committed
2
from abc import ABCMeta, abstractmethod
3

4
from mmcv.runner import BaseModule
wuyuefeng's avatar
wuyuefeng committed
5
6


7
class Base3DRoIHead(BaseModule, metaclass=ABCMeta):
zhangwenwei's avatar
zhangwenwei committed
8
    """Base class for 3d RoIHeads."""
wuyuefeng's avatar
wuyuefeng committed
9
10
11
12
13
14

    def __init__(self,
                 bbox_head=None,
                 mask_roi_extractor=None,
                 mask_head=None,
                 train_cfg=None,
15
16
17
18
                 test_cfg=None,
                 pretrained=None,
                 init_cfg=None):
        super(Base3DRoIHead, self).__init__(init_cfg=init_cfg)
wuyuefeng's avatar
wuyuefeng committed
19
20
21
22
23
24
25
26
27
28
29
30
31
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

        if bbox_head is not None:
            self.init_bbox_head(bbox_head)

        if mask_head is not None:
            self.init_mask_head(mask_roi_extractor, mask_head)

        self.init_assigner_sampler()

    @property
    def with_bbox(self):
zhangwenwei's avatar
zhangwenwei committed
32
        """bool: whether the RoIHead has box head"""
wuyuefeng's avatar
wuyuefeng committed
33
34
35
36
        return hasattr(self, 'bbox_head') and self.bbox_head is not None

    @property
    def with_mask(self):
zhangwenwei's avatar
zhangwenwei committed
37
        """bool: whether the RoIHead has mask head"""
wuyuefeng's avatar
wuyuefeng committed
38
39
40
41
        return hasattr(self, 'mask_head') and self.mask_head is not None

    @abstractmethod
    def init_bbox_head(self):
zhangwenwei's avatar
zhangwenwei committed
42
        """Initialize the box head."""
wuyuefeng's avatar
wuyuefeng committed
43
44
45
46
        pass

    @abstractmethod
    def init_mask_head(self):
zhangwenwei's avatar
zhangwenwei committed
47
        """Initialize maek head."""
wuyuefeng's avatar
wuyuefeng committed
48
49
50
51
        pass

    @abstractmethod
    def init_assigner_sampler(self):
zhangwenwei's avatar
zhangwenwei committed
52
        """Initialize assigner and sampler."""
wuyuefeng's avatar
wuyuefeng committed
53
54
55
56
57
        pass

    @abstractmethod
    def forward_train(self,
                      x,
zhangwenwei's avatar
zhangwenwei committed
58
                      img_metas,
wuyuefeng's avatar
wuyuefeng committed
59
60
61
62
63
                      proposal_list,
                      gt_bboxes,
                      gt_labels,
                      gt_bboxes_ignore=None,
                      **kwargs):
zhangwenwei's avatar
zhangwenwei committed
64
        """Forward function during training.
wuyuefeng's avatar
wuyuefeng committed
65
66
67
68
69

        Args:
            x (dict): Contains features from the first stage.
            img_metas (list[dict]): Meta info of each image.
            proposal_list (list[dict]): Proposal information from rpn.
zhangwenwei's avatar
zhangwenwei committed
70
            gt_bboxes (list[:obj:`BaseInstance3DBoxes`]):
wuyuefeng's avatar
wuyuefeng committed
71
72
                GT bboxes of each sample. The bboxes are encapsulated
                by 3D box structures.
73
74
            gt_labels (list[torch.LongTensor]): GT labels of each sample.
            gt_bboxes_ignore (list[torch.Tensor], optional):
zhangwenwei's avatar
zhangwenwei committed
75
                Ground truth boxes to be ignored.
wuyuefeng's avatar
wuyuefeng committed
76
77

        Returns:
wangtai's avatar
wangtai committed
78
            dict[str, torch.Tensor]: Losses from each head.
wuyuefeng's avatar
wuyuefeng committed
79
        """
wuyuefeng's avatar
wuyuefeng committed
80
81
82
83
84
        pass

    def simple_test(self,
                    x,
                    proposal_list,
zhangwenwei's avatar
zhangwenwei committed
85
                    img_metas,
wuyuefeng's avatar
wuyuefeng committed
86
87
88
89
90
91
92
93
94
                    proposals=None,
                    rescale=False,
                    **kwargs):
        """Test without augmentation."""
        pass

    def aug_test(self, x, proposal_list, img_metas, rescale=False, **kwargs):
        """Test with augmentations.

95
96
        If rescale is False, then returned bboxes and masks will fit the scale
        of imgs[0].
wuyuefeng's avatar
wuyuefeng committed
97
98
        """
        pass