Class EnsureChannelFirst

monai/transforms/utility/array.py:175–234

Adjust or add the channel dimension of input data to ensure `channel_first` shape. This extracts the `original_channel_dim` info from provided meta_data dictionary or MetaTensor input. This value should state which dimension is the channel dimension so that it can be moved forward, or

class EnsureChannelFirst(Transform):
    """
    Adjust or add the channel dimension of input data to ensure `channel_first` shape.

    This extracts the `original_channel_dim` info from provided meta_data dictionary or MetaTensor input. This value
    should state which dimension is the channel dimension so that it can be moved forward, or contain "no_channel" to
    state no dimension is the channel and so a 1-size first dimension is to be added.

    Args:
        strict_check: whether to raise an error when the meta information is insufficient.
        channel_dim: This argument can be used to specify the original channel dimension (integer) of the input array.
            It overrides the `original_channel_dim` from provided MetaTensor input.
            If the input array doesn't have a channel dim, this value should be ``'no_channel'``.
            If this is set to `None`, this class relies on `img` or `meta_dict` to provide the channel dimension.
    """

    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]

    def __init__(self, strict_check: bool = True, channel_dim: None | str | int = None):
        self.strict_check = strict_check
        self.input_channel_dim = channel_dim

    def __call__(self, img: torch.Tensor, meta_dict: Mapping | None = None) -> torch.Tensor:
        """
        Apply the transform to `img`.
        """
        if not isinstance(img, MetaTensor) and not isinstance(meta_dict, Mapping):
            if self.input_channel_dim is None:
                msg = "Metadata not available and channel_dim=None, EnsureChannelFirst is not in use."
                if self.strict_check:
                    raise ValueError(msg)
                warnings.warn(msg)
                return img
            else:
                img = MetaTensor(img)

        if isinstance(img, MetaTensor):
            meta_dict = img.meta

        channel_dim = meta_dict.get(MetaKeys.ORIGINAL_CHANNEL_DIM, None) if isinstance(meta_dict, Mapping) else None
        if self.input_channel_dim is not None:
            channel_dim = float("nan") if self.input_channel_dim == "no_channel" else self.input_channel_dim

        if channel_dim is None:
            msg = "Unknown original_channel_dim in the MetaTensor meta dict or `meta_dict` or `channel_dim`."
            if self.strict_check:
                raise ValueError(msg)
            warnings.warn(msg)
            return img

        # track the original channel dim
        if isinstance(meta_dict, dict):
            meta_dict[MetaKeys.ORIGINAL_CHANNEL_DIM] = channel_dim

        if is_no_channel(channel_dim):
            result = img[None]
        else:
            result = moveaxis(img, int(channel_dim), 0)  # type: ignore