26 lines
688 B
Python
26 lines
688 B
Python
from pathlib import Path
|
|
import pandas as pd
|
|
import os
|
|
|
|
from .version import Version
|
|
from .base import BaseDataset
|
|
from .info import DatasetInfo
|
|
from . import constants as C
|
|
|
|
|
|
def get_dataset(split="train"):
|
|
data_dir = Path(__file__).parent.parent.parent / "data"
|
|
# here = os.path.dirname(__path__)
|
|
# data_dir = here + '/../../data/'
|
|
|
|
meta = pd.read_csv(data_dir / "non-empty-meta.csv")
|
|
info = DatasetInfo.load(data_dir / "info.yaml")
|
|
|
|
# meta['image_path'] = [f'data/image/{panoid}_{heading}.jpg' for panoid, heading in zip(meta['panoid'], meta['heading'])]
|
|
|
|
ds = BaseDataset(info, meta)
|
|
if split:
|
|
return ds[split]
|
|
else:
|
|
return ds
|
|
|