from .segmentor3d_template import Segmentor3DTemplate
from .simple_seg import SimpleSeg

__all__ = {
    'Segmentor3DTemplate': Segmentor3DTemplate,
    'SimpleSeg': SimpleSeg,
}


def build_segmentor(model_cfg, runtime_cfg, dataset):
    model = __all__[model_cfg.NAME](
        model_cfg=model_cfg, runtime_cfg=runtime_cfg, dataset=dataset
    )

    return model
