Replace with pyfits with astropy imports