Spaces:
Running on Zero
Running on Zero
File size: 5,127 Bytes
0122a25 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Data connector for multi-sensor dataset."""
from __future__ import annotations
from typing import NamedTuple
from mapdet3d.data.typing import DictData, DictDataOrList
from .base import CallbackConnector, DataConnector, LossConnector
from .util import SourceKeyDescription, get_field_from_prediction
class MultiSensorDataConnector(DataConnector):
"""Data connector for multi-sensor data dict."""
def __init__(self, key_mapping: dict[str, str | SourceKeyDescription]):
"""Initializes the data connector with static remapping of the keys.
Args:
key_mapping (dict[str, | SourceKeyDescription]): Defines which
kwargs to pass onto the module.
TODO: Add Simple Example Configuration:
"""
_key_mapping = {}
multi_sensor_key_mapping = {}
for k, v in key_mapping.items():
if isinstance(v, dict):
sensors = v.get("sensors")
if sensors is not None:
multi_sensor_key_mapping[k] = v
else:
_key_mapping[k] = v["key"]
else:
_key_mapping[k] = v
super().__init__(_key_mapping)
self.multi_sensor_key_mapping = multi_sensor_key_mapping
def __call__(self, data: DictDataOrList) -> DictData:
"""Returns the train input for the model."""
input_dict = super().__call__(data)
for target_key, source_key in self.multi_sensor_key_mapping.items():
key = source_key["key"]
sensors = source_key["sensors"]
if isinstance(data, list):
input_dict[target_key] = [
[d[sensor][key] for sensor in sensors] for d in data
]
else:
input_dict[target_key] = [
data[sensor][key] for sensor in sensors
]
return input_dict
class MultiSensorLossConnector(LossConnector):
"""Multi-sensor Data connector for loss module of the training pipeline."""
def __call__(
self, prediction: DictData | NamedTuple, data: DictData
) -> DictData:
"""Returns the kwargs that are passed to the loss module.
Args:
prediction (DictData | NamedTuple): The output from model.
data (DictData): The data dictionary from the dataloader which
contains all data that was loaded.
Returns:
DictData: kwargs that are passed onto the loss.
"""
return get_multi_sensor_inputs(self.key_mapping, prediction, data)
class MultiSensorCallbackConnector(CallbackConnector):
"""Multi-sensor data connector for the callback."""
def __call__(
self, prediction: DictData | NamedTuple, data: DictData
) -> DictData:
"""Returns the kwargs that are passed to the callback.
Args:
prediction (DictData | NamedTuple): The output from model.
data (DictData): The data dictionary from the dataloader which
contains all data that was loaded.
Returns:
DictData: kwargs that are passed onto the callback.
"""
return get_multi_sensor_inputs(self.key_mapping, prediction, data)
def get_multi_sensor_inputs(
connection_dict: dict[str, SourceKeyDescription],
prediction: DictData | NamedTuple,
data: DictData,
) -> DictData:
"""Extracts multi-sensor input data from the provided SourceKeyDescription.
Args:
connection_dict (dict[str, SourceKeyDescription]): Input Key
description which is used to gather and remap data from the
two data dicts.
prediction (DictData): Dict containing the model prediction output.
data (DictData): Dict containing the dataloader output.
Raises:
ValueError: If the datasource is invalid.
Returns:
out (DictData): Dict containing new kwargs consisting of new key name
and data extracted from the data dicts.
"""
out: DictData = {}
for new_key_name, old_key_name in connection_dict.items():
# Assign field from data
if old_key_name["source"] == "data":
sensors = old_key_name.get("sensors")
if sensors is None:
if old_key_name["key"] not in data:
raise ValueError(
f"Key {old_key_name['key']} not found in data dict."
f" Available keys: {data.keys()}"
)
out[new_key_name] = data[old_key_name["key"]]
else:
out[new_key_name] = [
data[sensor][old_key_name["key"]] for sensor in sensors
]
# Assign field from prediction
elif old_key_name["source"] == "prediction":
out[new_key_name] = get_field_from_prediction(
prediction, old_key_name
)
else:
raise ValueError(
f"Unknown data source {old_key_name['source']}."
f"Available: [prediction, data]"
)
return out
|