spine.io.collate
Contains implementations of data collation classes.
Collate classes are a middleware between parsers and datasets. They are given
to torch.utils.data.DataLoader as the collate_fn argumement.
Classes
|
General collate function for all data types coming from the parsers. |
- class spine.io.collate.CollateAll(data_types: Mapping[str, str], split: bool = False, target_id: int = 0, source: Mapping[str, str] | None = None, overlay: Mapping[str, Any] | None = None, overlay_methods: Mapping[str, str] | None = None)[source]
General collate function for all data types coming from the parsers.
Provide it with a list of dictionaries. Each value can be one of:
A ParserTensor with coordinates, features and metadata, merged into rows of the form [batch_id, *coords, *features]
A feature-only ParserTensor, merged into [batch_id, *features]
A ParserIndex, ParserIndexList or ParserEdgeIndex, merged into an offset-adjusted index batch
Scalar values, lists and objects, gathered into a list
Methods
__call__(batch)Takes a list of parsed information, one per event in a batch, and collates them into a single object per entry in the batch.
stack_coord_tensors(batch, key)Stack coordinate tensors together across an overlay.
stack_feat_tensors(batch, key)Stack feature tensors together across an overlay.
stack_index_tensors(batch, key)Stack index tensors together across an overlay.
- name = 'all'
- stack_coord_tensors(batch: Sequence[dict[str, Any]], key: str) TensorBatch[source]
Stack coordinate tensors together across an overlay.
- Parameters:
batch (List[Dict]) – List of dictionaries of parsed information, one per event. Each dictionary matches one data key to one event-worth of parsed data.
key (str) – Data product key
- Returns:
Batched coordinate tensor
- Return type:
- stack_index_tensors(batch: Sequence[dict[str, Any]], key: str) IndexBatch | EdgeIndexBatch[source]
Stack index tensors together across an overlay.
- Parameters:
batch (List[Dict]) – List of dictionaries of parsed information, one per event. Each dictionary matches one data key to one event-worth of parsed data.
key (str) – Data product key
- Returns:
Batched index tensor
- Return type:
Union[IndexBatch, EdgeIndexBatch]
- stack_feat_tensors(batch: Sequence[dict[str, Any]], key: str) TensorBatch[source]
Stack feature tensors together across an overlay.
- Parameters:
batch (List[Dict]) – List of dictionaries of parsed information, one per event. Each dictionary matches one data key to one event-worth of parsed data.
key (str) – Data product key
- Returns:
Batched feature tensor
- Return type: