Repository: Project-MONAI/MONAI
File: monai/data/image_writer.py (Python)

Code before:
    .. code-block:: python

        import numpy as np
        from monai.data import ITKWriter

        np_data = np.arange(48).reshape(3, 4, 4)

        # write as 3d spatial image no channel
        writer = ITKWriter(output_dtype=np.float32)
        writer.set_data_array(np_data, channel_dim=None)
        # optionally set metadata affine
        writer.set_metadata({"affine": np.eye(4), "original_affine": -1 * np.eye(4)})
        writer.write("test1.nii.gz")

        # write as 2d image, channel-first
        writer = ITKWriter(output_dtype=np.uint8)
        writer.set_data_array(np_data, channel_dim=0)
        writer.set_metadata({"spatial_shape": (5, 5)})
        writer.write("test1.png")

    """

    output_dtype: DtypeLike = None
    channel_dim: int | None

    def __init__(self, output_dtype: DtypeLike = np.float32, affine_lps_to_ras: bool | None = None, **kwargs):
        """
        Args:
            output_dtype: output data type.
            affine_lps_to_ras: whether to convert the affine matrix from "LPS" to "RAS". Defaults to ``None``.
                Set to ``True`` to be consistent with ``NibabelWriter``,
                otherwise the affine matrix is assumed already in the ITK convention.
            kwargs: keyword arguments passed to ``ImageWriter``.

        The constructor will create ``self.output_dtype`` internally.
        ``affine`` and ``channel_dim`` are initialized as instance members (default ``None``, ``0``):

            - user-specified ``affine`` should be set in ``set_metadata``,
            - user-specified ``channel_dim`` should be set in ``set_data_array``.
        """
        super().__init__(
            output_dtype=output_dtype, affine_lps_to_ras=affine_lps_to_ras, affine=None, channel_dim=0, **kwargs
        )

    def set_data_array(
        self, data_array: NdarrayOrTensor, channel_dim: int | None = 0, squeeze_end_dims: bool = True, **kwargs
    ):
        """
        Convert ``data_array`` into 'channel-last' numpy ndarray.

        Args:

Review comment (suggestion):
Not sure why we need this default value change here? But I find another wired thing just using the data provided in this [issue](https://github.com/Project-MONAI/MONAI/issues/6136). The "nrrd" data after `SaveImage` will have affine with shape (5, 5), It looks like it assumes that the data becomes 4-dimensional, without excluding the channel dimension.

```
original_data = LoadImage(image_only=True)(original_file)
SaveImage(output_dir="./monai_modified_meta",
          output_ext = ".nrrd",
          writer= "ITKWriter",
          resample=False,
          separate_folder= False,
          squeeze_end_dims = False,
          output_dtype=itk.SI)(original_data)
print(original_data.meta["spatial_shape"]) # --> ([512 512   1])

save_data = LoadImage(image_only=True)(save_data_path)
print(save_data.meta["spatial_shape"]) # --> [512 512   1   1]
```

Code after:
    .. code-block:: python

        import numpy as np
        from monai.data import ITKWriter

        np_data = np.arange(48).reshape(3, 4, 4)

        # write as 3d spatial image no channel
        writer = ITKWriter(output_dtype=np.float32)
        writer.set_data_array(np_data, channel_dim=None)
        # optionally set metadata affine
        writer.set_metadata({"affine": np.eye(4), "original_affine": -1 * np.eye(4)})
        writer.write("test1.nii.gz")

        # write as 2d image, channel-first
        writer = ITKWriter(output_dtype=np.uint8)
        writer.set_data_array(np_data, channel_dim=0)
        writer.set_metadata({"spatial_shape": (5, 5)})
        writer.write("test1.png")

    """

    output_dtype: DtypeLike = None
    channel_dim: int | None

    def __init__(self, output_dtype: DtypeLike = np.float32, affine_lps_to_ras: bool | None = True, **kwargs):
        """
        Args:
            output_dtype: output data type.
            affine_lps_to_ras: whether to convert the affine matrix from "LPS" to "RAS". Defaults to ``True``.
                Set to ``True`` to be consistent with ``NibabelWriter``,
                otherwise the affine matrix is assumed already in the ITK convention.
                Set to ``None`` to use ``data_array.meta[MetaKeys.SPACE]`` to determine the flag.
            kwargs: keyword arguments passed to ``ImageWriter``.

        The constructor will create ``self.output_dtype`` internally.
        ``affine`` and ``channel_dim`` are initialized as instance members (default ``None``, ``0``):

            - user-specified ``affine`` should be set in ``set_metadata``,
            - user-specified ``channel_dim`` should be set in ``set_data_array``.
        """
        super().__init__(
            output_dtype=output_dtype, affine_lps_to_ras=affine_lps_to_ras, affine=None, channel_dim=0, **kwargs
        )

    def set_data_array(
        self, data_array: NdarrayOrTensor, channel_dim: int | None = 0, squeeze_end_dims: bool = True, **kwargs
    ):
        """
        Convert ``data_array`` into 'channel-last' numpy ndarray.