Repository: Project-MONAI/MONAI
File: monai/transforms/spatial/dictionary.py (Python)

Code before:
                in this dictionary must be tensor like arrays that are channel first and have at most
                three spatial dimensions
            lazy: a flag to indicate whether this transform should execute lazily or not
                during this call. Setting this to False or True overrides the ``lazy`` flag set
                during initialization for this call. Defaults to None.

        Returns:
            a dictionary containing the transformed data, as well as any other data present in the dictionary
        """
        d: dict = dict(data)

        _init_shape, _pixdim, should_match = None, None, False
        output_shape_k = None  # tracking output shape
        lazy_ = self.lazy if lazy is None else lazy

        for key, mode, padding_mode, align_corners, dtype, scale_extent in self.key_iterator(
            d, self.mode, self.padding_mode, self.align_corners, self.dtype, self.scale_extent
        ):
            if self.ensure_same_shape and isinstance(d[key], MetaTensor):
                if _init_shape is None and _pixdim is None:
                    _init_shape, _pixdim = d[key].peek_pending_shape(), d[key].pixdim
                else:
                    should_match = np.allclose(_init_shape, d[key].peek_pending_shape()) and np.allclose(
                        _pixdim, d[key].pixdim, atol=1e-3
                    )
                    _pixdim = d[key].pixdim
            d[key] = self.spacing_transform(
                data_array=d[key],
                mode=mode,
                padding_mode=padding_mode,
                align_corners=align_corners,
                dtype=dtype,
                scale_extent=scale_extent,
                output_spatial_shape=output_shape_k if should_match else None,
                lazy=lazy_,
            )
            output_shape_k = d[key].peek_pending_shape() if isinstance(d[key], MetaTensor) else d[key].shape[1:]
        return d

    def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.spacing_transform.inverse(cast(torch.Tensor, d[key]))
        return d


class Orientationd(MapTransform, InvertibleTransform, LazyTransform):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.Orientation`.

    This transform assumes the channel-first input format.

Review comment (suggestion):
This whole if-else block could be:

```python
                if _init_shape is None:
                    _init_shape = d[key].peek_pending_shape()
                else:
                    should_match = np.allclose(_init_shape, d[key].peek_pending_shape()) and np.allclose(
                        _pixdim, d[key].pixdim, atol=1e-3
                    )
                _pixdim = d[key].pixdim
```

Code after:
                in this dictionary must be tensor like arrays that are channel first and have at most
                three spatial dimensions
            lazy: a flag to indicate whether this transform should execute lazily or not
                during this call. Setting this to False or True overrides the ``lazy`` flag set
                during initialization for this call. Defaults to None.

        Returns:
            a dictionary containing the transformed data, as well as any other data present in the dictionary
        """
        d: dict = dict(data)

        _init_shape, _pixdim, should_match = None, None, False
        output_shape_k = None  # tracking output shape
        lazy_ = self.lazy if lazy is None else lazy

        for key, mode, padding_mode, align_corners, dtype, scale_extent in self.key_iterator(
            d, self.mode, self.padding_mode, self.align_corners, self.dtype, self.scale_extent
        ):
            if self.ensure_same_shape and isinstance(d[key], MetaTensor):
                if _init_shape is None and _pixdim is None:
                    _init_shape, _pixdim = d[key].peek_pending_shape(), d[key].pixdim
                else:
                    should_match = np.allclose(_init_shape, d[key].peek_pending_shape()) and np.allclose(
                        _pixdim, d[key].pixdim, atol=1e-3
                    )
            d[key] = self.spacing_transform(
                data_array=d[key],
                mode=mode,
                padding_mode=padding_mode,
                align_corners=align_corners,
                dtype=dtype,
                scale_extent=scale_extent,
                output_spatial_shape=output_shape_k if should_match else None,
                lazy=lazy_,
            )
            if output_shape_k is None:
                output_shape_k = d[key].peek_pending_shape() if isinstance(d[key], MetaTensor) else d[key].shape[1:]
        return d

    def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.spacing_transform.inverse(cast(torch.Tensor, d[key]))
        return d


class Orientationd(MapTransform, InvertibleTransform, LazyTransform):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.Orientation`.

    This transform assumes the channel-first input format.