data

Data loading: standard sample keys, packed-batch collation, and the point cloud data loader.

Classes:

DataKeys –

Standard keys of a sample dict, shared by the datasets, transforms and models.
PointCloudDataLoader –

DataLoader that batches point clouds with the packed-batch collate by default.

Functions:

collate –

Collate a list of point-cloud sample dicts into one batched dict.

DataKeys ¶

Bases: StrEnum

Standard keys of a sample dict, shared by the datasets, transforms and models.

Members are plain strings, so a key is usable wherever a literal is (e.g. data[DataKeys.POS] and data["pos"] address the same entry).

PointCloudDataLoader ¶

PointCloudDataLoader(
    dataset: Dataset,
    *,
    batch_from: str = POS,
    batch_key: str = BATCH,
    stack_keys: Optional[Sequence[str]] = None,
    cat_keys: Optional[Sequence[str]] = None,
    **kwargs: Any,
)

Bases: DataLoader

DataLoader that batches point clouds with the packed-batch collate by default.

Wraps torch.utils.data.DataLoader, defaulting collate_fn to collate. How keys collate is set by the spec arguments (batch_from / batch_key for the per-point index, stack_keys for dense per-scene ground truth, cat_keys for ragged per-scene ground truth). These are supplied by the caller, never read off the dataset: transforms rewrite the key set downstream of the dataset (a box key may be derived from an object by a transform), so only the code building the loader knows which keys must stack or cat. Passing collate_fn=... via the usual DataLoader kwarg overrides the spec.

Parameters:

dataset (Dataset) –

The dataset to load from.
batch_from (str, default: POS ) –

Key whose leading dimension defines the per-point batch index.
batch_key (str, default: BATCH ) –

Output key for the per-point batch index.
stack_keys (Optional[Sequence[str]], default: None ) –

Keys collated by stacking to a leading batch dim instead of concatenating.
cat_keys (Optional[Sequence[str]], default: None ) –

Packed keys that additionally emit a batch_<key> per-element scene index.
**kwargs (Any, default: {} ) –

Forwarded to torch.utils.data.DataLoader (batch_size, shuffle, collate_fn, ...).

collate ¶

collate(
    data_list: List[Dict[str, Any]],
    batch_from: str = POS,
    batch_key: str = BATCH,
    stack_keys: Optional[KeyCollection] = None,
    cat_keys: Optional[KeyCollection] = None,
) -> Dict[str, Any]

Collate a list of point-cloud sample dicts into one batched dict.

By default every per-point tensor is concatenated PyG-style along dim 0 (packed), scalars are stacked, and a per-point batch_key index is synthesized from batch_from. Two extra knobs say how specific keys collate instead:

stack_keys: stack to a new leading batch dim (\((M, \cdot) \to (B, M, \cdot)\), \((N, \cdot) \to (B, N, \cdot)\)) rather than concatenating. Used for fixed-size per-scene ground truth (the VoteNet loss consumes dense \((B, M, \cdot)\) targets, which a plain cat would flatten).
cat_keys: keep these packed (cat) but additionally emit a batch_<key> scene index mirroring batch_key. Used for ragged per-scene ground truth such as box \((K, 8)\) -> batch_box \((K,)\).

Every key must be present in every sample; a key missing from a sample raises a ValueError. stack_keys / cat_keys entries absent from all samples are ignored. A key may appear in only one of stack_keys / cat_keys; overlapping entries raise a ValueError.

Parameters:

data_list (List[Dict[str, Any]]) –

List of sample dicts.
batch_from (str, default: POS ) –

Key whose leading dimension defines the per-point batch index.
batch_key (str, default: BATCH ) –

Output key for the per-point batch index.
stack_keys (Optional[KeyCollection], default: None ) –

Keys collated by stacking to a leading batch dim instead of concatenating.
cat_keys (Optional[KeyCollection], default: None ) –

Packed keys that additionally emit a batch_<key> per-element scene index.

Returns:

Dict[str, Any] –

A single batched dict.