import numpy as np
import nibabel as nib
from monai.transforms import SaveImaged, LoadImaged
from monai.data import ITKWriter, NibabelWriter

# create a testing volume
img = nib.Nifti1Image(np.ones((3, 4, 5, 1, 1)), np.eye(4))
img.header.set_qform(np.eye(4), code=2)
nib.save(img, "test_qform.nii.gz")

# load/save test
data_in = LoadImaged(keys="img", reader="nibabelreader")({"img": "test_qform.nii.gz"})
SaveImaged(keys="img", output_postfix="nib", writer=NibabelWriter)(data_in)
SaveImaged(keys="img", output_postfix="itk", writer=ITKWriter)(data_in)

# print the saved
ori_hdr = nib.load("test_qform.nii.gz").header
print(ori_hdr['sform_code'], ori_hdr['qform_code'])
itk_hdr = nib.load("test_qform/test_qform_itk.nii.gz").header
print(itk_hdr['sform_code'], itk_hdr['qform_code'])
nib_hdr = nib.load("test_qform/test_qform_nib.nii.gz").header
print(nib_hdr['sform_code'], nib_hdr['qform_code'])