File size: 1,687 Bytes
f71ac1d | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # pylint: disable=no-member
"""Select Sensor transformation."""
from vis4d.data.typing import DictData
from .base import Transform
@Transform("data", "data")
class SelectSensor:
"""Keep data from one sensor only but keep shared data.
Note: The input data is assumed to be in the format of DictData[DictData],
i.e. a list of data dictionaries, each of which contains a dictionary of
either the data from a sensor or the shared data (metadata) for all
sensors.
Example:
>>> data = [
{
"sensor1": {"image": 1, "label": 2},
"sensor2": {"image": 1, "label": 2},
"meta": 3},
},
]
>>> tsfm = SelectSensor(
sensor="sensor1", sensors=["sensor1", "sensor2"]
)
>>> tsfm(data)
[{"image": 1, "label": 2, "meta": 3},]
"""
def __init__(self, selected_sensor: str) -> None:
"""Creates an instance of SelectSensor.
Args:
selected_sensor (str): The name of the sensor to keep.
"""
self.selected_sensor = selected_sensor
def __call__(self, batch: list[DictData]) -> list[DictData]:
"""Select data from one sensor only."""
output_batch = []
for data in batch:
output_data = {}
for key in data.keys():
if key in self.sensors: # type: ignore
if key == self.selected_sensor:
output_data.update(data[key])
else:
output_data[key] = data[key]
output_batch.append(output_data)
return output_batch
|