Cube to image

def cube_to_image(cube, slicepos=None):
    """ Make an image out of a cube.
    Both in- and output should by fits.HDUs"""
    header = cube.header.copy()
    header['NAXIS'] = 2
    del header['NAXIS3']
    del header['CRVAL3']
    del header['CDELT3']
    del header['CTYPE3']
    del header['CRPIX3']
    del header['CUNIT3']
    if slicepos:
        data = cube.data[slicepos]
    else:
        data = cube.data.sum(0)
    return fits.ImageHDU(data, header)