base.py 5.1 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
jshilong's avatar
jshilong committed
2
3
from typing import List, Union

4
from mmdet3d.core import Det3DDataSample
5
6
from mmdet3d.core.utils import (ForwardResults, InstanceList, OptConfigType,
                                OptMultiConfig, OptSampleList, SampleList)
7
from mmdet3d.registry import MODELS
8
from mmdet.models import BaseDetector
zhangwenwei's avatar
zhangwenwei committed
9
10


11
@MODELS.register_module()
zhangwenwei's avatar
zhangwenwei committed
12
class Base3DDetector(BaseDetector):
13
    """Base class for 3D detectors.
zhangwenwei's avatar
zhangwenwei committed
14

15
    Args:
16
17
18
19
20
       data_preprocessor (dict or ConfigDict, optional): The pre-process
           config of :class:`BaseDataPreprocessor`.  it usually includes,
            ``pad_size_divisor``, ``pad_value``, ``mean`` and ``std``.
       init_cfg (dict or ConfigDict, optional): the config to control the
           initialization. Defaults to None.
21
22
23
    """

    def __init__(self,
24
25
26
                 data_processor: OptConfigType = None,
                 init_cfg: OptMultiConfig = None) -> None:
        super().__init__(data_preprocessor=data_processor, init_cfg=init_cfg)
27

28
    def forward(self,
jshilong's avatar
jshilong committed
29
30
                inputs: Union[dict, List[dict]],
                data_samples: OptSampleList = None,
31
32
33
                mode: str = 'tensor',
                **kwargs) -> ForwardResults:
        """The unified entry for a forward process in both training and test.
34

35
        The method should accept three modes: "tensor", "predict" and "loss":
36

37
38
39
40
41
42
        - "tensor": Forward the whole network and return tensor or tuple of
        tensor without any post-processing, same as a common nn.Module.
        - "predict": Forward and return the predictions, which are fully
        processed to a list of :obj:`DetDataSample`.
        - "loss": Forward and return a dict of losses according to the given
        inputs and data samples.
43

44
45
        Note that this method doesn't handle neither back propagation nor
        optimizer updating, which are done in the :meth:`train_step`.
46
47

        Args:
jshilong's avatar
jshilong committed
48
49
50
51
52
53
54
55
56
57
58
59
60
            inputs  (dict | list[dict]): When it is a list[dict], the
                outer list indicate the test time augmentation. Each
                dict contains batch inputs
                which include 'points' and 'imgs' keys.

                - points (list[torch.Tensor]): Point cloud of each sample.
                - imgs (torch.Tensor): Image tensor has shape (B, C, H, W).
            data_samples (list[:obj:`DetDataSample`],
                list[list[:obj:`DetDataSample`]], optional): The
                annotation data of every samples. When it is a list[list], the
                outer list indicate the test time augmentation, and the
                inter list indicate the batch. Otherwise, the list simply
                indicate the batch. Defaults to None.
61
            mode (str): Return what kind of value. Defaults to 'tensor'.
62
63

        Returns:
64
            The return type depends on ``mode``.
65

66
67
68
69
70
            - If ``mode="tensor"``, return a tensor or a tuple of tensor.
            - If ``mode="predict"``, return a list of :obj:`DetDataSample`.
            - If ``mode="loss"``, return a dict of tensor.
        """
        if mode == 'loss':
jshilong's avatar
jshilong committed
71
            return self.loss(inputs, data_samples, **kwargs)
72
        elif mode == 'predict':
jshilong's avatar
jshilong committed
73
74
75
76
77
78
79
80
81
82
            if isinstance(data_samples[0], list):
                # aug test
                assert len(data_samples[0]) == 1, 'Only support ' \
                                                  'batch_size 1 ' \
                                                  'in mmdet3d when ' \
                                                  'do the test' \
                                                  'time augmentation.'
                return self.aug_test(inputs, data_samples, **kwargs)
            else:
                return self.predict(inputs, data_samples, **kwargs)
83
        elif mode == 'tensor':
jshilong's avatar
jshilong committed
84
            return self._forward(inputs, data_samples, **kwargs)
85
        else:
86
87
            raise RuntimeError(f'Invalid mode "{mode}". '
                               'Only supports loss, predict and tensor mode')
88

89
90
    def convert_to_datasample(self, results_list: InstanceList) -> SampleList:
        """Convert results list to `Det3DDataSample`.
91

92
93
        Subclasses could override it to be compatible for some multi-modality
        3D detectors.
94
95
96
97
98
99
100

        Args:
            results_list (list[:obj:`InstanceData`]): Detection results of
                each sample.

        Returns:
            list[:obj:`Det3DDataSample`]: Detection results of the
101
102
            input. Each Det3DDataSample usually contains
            'pred_instances_3d'. And the ``pred_instances_3d`` usually
103
104
105
            contains following keys.

                - scores_3d (Tensor): Classification scores, has a shape
106
107
                    (num_instance, )
                - labels_3d (Tensor): Labels of 3D bboxes, has a shape
108
                    (num_instances, ).
109
110
111
112
                - bboxes_3d (Tensor): Contains a tensor with shape
                    (num_instances, C) where C >=7.
        """
        out_results_list = []
113
114
115
        for i in range(len(results_list)):
            result = Det3DDataSample()
            result.pred_instances_3d = results_list[i]
116
117
            out_results_list.append(result)
        return out_results_list