from __future__ import annotations
from typing import Any, Callable, Dict, Optional
import torch
from torch.utils.data import Dataset
from ..index import ScenarioIndex
from ._base import PathLike, check_parts_for, ensure_parts_for, required_parts, resolve_index
from .anomaly_lidar import AnomalyLiDARDataset
from .pointcloud import PointCloudDataset
[docs]
class LiDARDataset(Dataset):
"""Composite dataset for LiDAR point clouds with anomaly labels.
Wraps :class:`PointCloudDataset` and :class:`AnomalyLiDARDataset`.
Returns a dict with keys ``'points'`` (``List[DataFrame]``) and
``'anomaly_lidar'`` (``List[BoolTensor]``), plus the evaluator
identifiers ``'scenario_id'`` and ``'timesteps'``.
Parameters
----------
root:
Path to the CarlAnomaly dataset directory.
split:
``'train'``, ``'test_normal'``, ``'test_anomaly'``, or ``'test'``.
transform:
Optional transform applied to the merged output dict.
download:
If ``True``, fetch the archive parts this composite needs (``base`` +
``lidar``) into ``root`` before loading.
By default (``False``) the parts must already be there; a missing
part raises :class:`FileNotFoundError` with download instructions.
Additional keyword arguments (``clip_len``, ``stride``, ``parts``, ...)
are forwarded to :class:`~carlanomaly.index.ScenarioIndex`.
"""
def __init__(
self,
root: Optional[PathLike] = None,
split: str = "train",
*,
transform: Optional[Callable] = None,
index: Optional[ScenarioIndex] = None,
download: bool = False,
**index_kwargs: Any,
) -> None:
parts = index_kwargs.pop("parts", None) or required_parts(
[("pointcloud", None), ("anomaly_lidar", None)]
)
if download:
ensure_parts_for(root, split, index, parts,
verify=index_kwargs.get("download_verify", True))
elif index is None:
check_parts_for(root, split, parts)
index = resolve_index(root, split, index=index, download=False, **index_kwargs)
self._index = index
self.transform = transform
# Parts already fetched above; children must not re-download.
self.pointcloud = PointCloudDataset(index=index, download=False)
self.anomaly_lidar = AnomalyLiDARDataset(index=index, download=False)
@property
def index(self) -> ScenarioIndex:
"""The :class:`ScenarioIndex` this dataset is built on (shareable)."""
return self._index
[docs]
def __len__(self) -> int:
return len(self._index)
[docs]
def __getitem__(self, idx: int) -> Dict:
rec, _ = self._index[idx]
item = {
"scenario_id": str(rec.path),
"timesteps": torch.tensor(self._index.timesteps_for(idx), dtype=torch.long),
"points": self.pointcloud[idx],
"anomaly_lidar": self.anomaly_lidar[idx],
}
if self.transform is not None:
item = self.transform(item)
return item