File size: 808 Bytes
4187c6f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from ..base import DataBase
from .dataset import BEVKitti360Dataset
from ..schema import KITTIDataConfiguration
class BEVKitti360Data(DataBase):
def __init__(self, cfg: KITTIDataConfiguration) -> None:
self.cfg = cfg
self._dataset = {}
def prepare_data(self) -> None:
return
def setup(self, stage: str) -> None:
split = {
'fit': 'train',
'val': 'val',
'validate': 'val',
'test': 'val',
"train": "train"
}[stage]
self._dataset[stage] = BEVKitti360Dataset(
cfg=self.cfg,
split_name=split
)
def dataset(self, stage: str):
if self._dataset.get(stage) is None:
self.setup(stage)
return self._dataset[stage]
|