Embed dicom metadata into a nifti file

embed_into_nifti.py

from copy import deepcopy

from dcmstack import DcmMetaExtension
import nibabel as nb

def embed_nifti(dcmfiles, niftifile, force=False):
    stack = ds.parse_and_stack(dcmfiles, force=force).values()
    if len(stack) > 1:
        raise ValueError('Found multiple series')
    stack = stack[0]
        
    #Create the nifti image using the data array
    if not os.path.exists(niftifile):
        nifti_image = stack.to_nifti(embed_meta=True)
        nifti_image.to_filename(niftifile)
        return
    
    nifti_image = nb.load(niftifile)
    nifti_header = nifti_image.get_header()

    data = stack.get_data()
    affine = stack.get_affine()
    
    #Figure out the number of three (or two) dimensional volumes
    n_vols = 1
    if len(data.shape) > 3:
        n_vols *= data.shape[3]
    if len(data.shape) > 4:
        n_vols *= data.shape[4]
    
    files_per_vol = len(stack._files_info) / n_vols
    slice_dim = 2
    n_slices = data.shape[slice_dim]
    
    #Build meta data for each volume if needed
    vol_meta = []
    if files_per_vol > 1:
        for vol_idx in xrange(n_vols):
            start_slice = vol_idx * n_slices
            end_slice = start_slice + n_slices
            exts = [file_info[0].meta_ext
                    for file_info in stack._files_info[start_slice:end_slice]]
            meta = DcmMetaExtension.from_sequence(exts, 2)
            vol_meta.append(meta)
    else:
        vol_meta = [file_info[0].meta_ext 
                    for file_info in stack._files_info]
                    
    #Build meta data for each time point / vector component
    if len(data.shape) == 5:
        if data.shape[3] != 1:
            vec_meta = []
            for vec_idx in xrange(data.shape[4]):
                start_idx = vec_idx * data.shape[3]
                end_idx = start_idx + data.shape[3]
                meta = DcmMetaExtension.from_sequence(\
                    vol_meta[start_idx:end_idx], 3)
                vec_meta.append(meta)
        else:
            vec_meta = vol_meta
                
        meta_ext = DcmMetaExtension.from_sequence(vec_meta, 4)
    elif len(data.shape) == 4:
        meta_ext = DcmMetaExtension.from_sequence(vol_meta, 3)
    else:
        meta_ext = vol_meta[0]
        if meta_ext is file_info[0].meta_ext:
            meta_ext = deepcopy(meta_ext)
            
    meta_ext.shape = data.shape
    meta_ext.slice_dim = slice_dim
    meta_ext.affine = nifti_header.get_best_affine()
    meta_ext.reorient_transform = np.eye(4)
            
    #Filter and embed the meta data
    meta_ext.filter_meta(stack._meta_filter)
    nifti_header.extensions = nb.nifti1.Nifti1Extensions([meta_ext])

    nifti_image.update_header()
    nifti_image.to_filename(niftifile)