__init__.py 672 Bytes
Newer Older
1
2
from .detector3d_template import Detector3DTemplate
from .PartA2_net import PartA2Net
3
from .point_rcnn import PointRCNN
Shaoshuai Shi's avatar
Shaoshuai Shi committed
4
5
6
from .pointpillar import PointPillar
from .pv_rcnn import PVRCNN
from .second_net import SECONDNet
7
from .second_net_iou import SECONDNetIoU
8
9
10
11
12

__all__ = {
    'Detector3DTemplate': Detector3DTemplate,
    'SECONDNet': SECONDNet,
    'PartA2Net': PartA2Net,
Gus-Guo's avatar
Gus-Guo committed
13
    'PVRCNN': PVRCNN,
14
    'PointPillar': PointPillar,
15
16
    'PointRCNN': PointRCNN,
    'SECONDNetIoU': SECONDNetIoU
17
18
19
20
21
22
23
24
25
}


def build_detector(model_cfg, num_class, dataset):
    model = __all__[model_cfg.NAME](
        model_cfg=model_cfg, num_class=num_class, dataset=dataset
    )

    return model