Ignore invalid keys instead of raising an error

kornia/augmentation/container/augment.py

@@ -175,8 +175,9 @@ class AugmentationSequential(TransformMatrixMinIn, ImageSequential):
  ... )
  >>> out = aug_list(input, mask, bbox)
 
- How to use a dictionary as input with AugmentationSequential? The dictionary should starts with
- one of the datakey availables.
+ How to use a dictionary as input with AugmentationSequential? The dictionary keys that start with
+ one of the available datakeys will be augmented accordingly. Otherwise, the dictionary item is passed
+ without any augmentation.
 
  >>> import kornia.augmentation as K
  >>> img = torch.randn(1, 3, 256, 256)
@@ -291,7 +292,7 @@ def inverse( # type: ignore[override]
  """
  original_keys = None
  if len(args) == 1 and isinstance(args[0], dict):
- original_keys, data_keys, args = self._preproc_dict_data(args[0])
+ original_keys, data_keys, args, invalid_data = self._preproc_dict_data(args[0])
 
  # args here should already be `DataType`
  # NOTE: how to right type to: unpacked args <-> tuple of args to unpack
@@ -324,7 +325,10 @@ def inverse( # type: ignore[override]
  outputs = self._arguments_postproc(args, outputs, data_keys=self.transform_op.data_keys) # type: ignore
 
  if isinstance(original_keys, tuple):
- return {k: v for v, k in zip(outputs, original_keys)}
+ result = {k: v for v, k in zip(outputs, original_keys)}
+ if invalid_data:
+ result.update(invalid_data)
+ return result
 
  if len(outputs) == 1 and isinstance(outputs, list):
  return outputs[0]
@@ -414,7 +418,7 @@ def forward( # type: ignore[override]
  # Unpack/handle dictionary args
  original_keys = None
  if len(args) == 1 and isinstance(args[0], dict):
- original_keys, data_keys, args = self._preproc_dict_data(args[0])
+ original_keys, data_keys, args, invalid_data = self._preproc_dict_data(args[0])
 
  self.transform_op.data_keys = self.transform_op.preproc_datakeys(data_keys)
 
@@ -455,7 +459,10 @@ def forward( # type: ignore[override]
  self._params = params
 
  if isinstance(original_keys, tuple):
- return {k: v for v, k in zip(outputs, original_keys)}
+ result = {k: v for v, k in zip(outputs, original_keys)}
+ if invalid_data:
+ result.update(invalid_data)
+ return result
 
  if len(outputs) == 1 and isinstance(outputs, list):
  return outputs[0]
@@ -464,17 +471,19 @@ def forward( # type: ignore[override]
 
  def _preproc_dict_data(
  self, data: Dict[str, DataType]
- ) -> Tuple[Tuple[str, ...], List[DataKey], Tuple[DataType, ...]]:
+ ) -> Tuple[Tuple[str, ...], List[DataKey], Tuple[DataType, ...], Optional[Dict[str, Any]]]:
  if self.data_keys is not None:
  raise ValueError("If you are using a dictionary as input, the data_keys should be None.")
 
  keys = tuple(data.keys())
- data_keys = self._read_datakeys_from_dict(keys)
+ data_keys, invalid_keys = self._read_datakeys_from_dict(keys)
+ invalid_data = {i: data.pop(i) for i in invalid_keys} if invalid_keys else None
+ keys = tuple(k for k in keys if k not in invalid_keys) if invalid_keys else keys
  data_unpacked = tuple(data.values())
 
- return keys, data_keys, data_unpacked
+ return keys, data_keys, data_unpacked, invalid_data
 
- def _read_datakeys_from_dict(self, keys: Sequence[str]) -> List[DataKey]:
+ def _read_datakeys_from_dict(self, keys: Sequence[str]) -> Tuple[List[DataKey], Optional[List[str]]]:
  def retrieve_key(key: str) -> DataKey:
  """Try to retrieve the datakey value by matching `<datakey>*`"""
  # Alias cases, like INPUT, will not be get by the enum iterator.
@@ -492,7 +501,15 @@ def retrieve_key(key: str) -> DataKey:
  f"Your input data dictionary keys should start with some of datakey values: {allowed_dk}. Got `{key}`"
  )
 
- return [DataKey.get(retrieve_key(k)) for k in keys]
+ valid_data_keys = []
+ invalid_keys = []
+ for k in keys:
+ try:
+ valid_data_keys.append(DataKey.get(retrieve_key(k)))
+ except ValueError:
+ invalid_keys.append(k)
+
+ return valid_data_keys, invalid_keys
 
  def _preproc_mask(self, arg: MaskDataType) -> MaskDataType:
  if isinstance(arg, list):

tests/augmentation/test_container.py

@@ -776,20 +776,22 @@ def test_dict_as_input_forward_and_inverse(self, random_apply, bbox_key, device,
  random_apply=random_apply,
  )
 
- data = {"input": inp, "mask": mask, bbox_key: bbox, "keypoints": keypoints}
+ data = {"input": inp, "mask": mask, bbox_key: bbox, "keypoints": keypoints, "id": 45}
  out = aug(data)
  assert out["input"].shape == inp.shape
  assert out["mask"].shape == mask.shape
  assert out[bbox_key].shape == bbox.shape
  assert out["keypoints"].shape == keypoints.shape
  assert set(out["mask"].unique().tolist()).issubset(set(mask.unique().tolist()))
+ assert out["id"] == 45
 
  out_inv = aug.inverse(out)
  assert out_inv["input"].shape == inp.shape
  assert out_inv["mask"].shape == mask.shape
  assert out_inv[bbox_key].shape == bbox.shape
  assert out_inv["keypoints"].shape == keypoints.shape
  assert set(out_inv["mask"].unique().tolist()).issubset(set(mask.unique().tolist()))
+ assert out_inv["id"] == 45
 
  if random_apply is False:
  reproducibility_test(data, aug)