diff options
-rwxr-xr-x | astro/fits/crop_image.py | 35 | ||||
-rwxr-xr-x | astro/fits/rescale_image.py | 183 |
2 files changed, 218 insertions, 0 deletions
diff --git a/astro/fits/crop_image.py b/astro/fits/crop_image.py new file mode 100755 index 0000000..9cfa7bf --- /dev/null +++ b/astro/fits/crop_image.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2017 Aaron LI <aly@aaronly.me> +# MIT license +# + +""" +Crop out the central region of specified size from the FITS image. +""" + +import argparse + +from rescale_image import FITSImage + + +def main(): + parser = argparse.ArgumentParser( + description="Rescale a FITS image to the desired size/resolution") + parser.add_argument("-C", "--clobber", action="store_true", + help="overwrite existing output file") + parser.add_argument("-s", "--size", type=float, required=True, + help="central crop box size [deg]") + parser.add_argument("-i", "--infile", required=True, + help="input FITS image") + parser.add_argument("-o", "--outfile", required=True, + help="output cropped FITS image") + args = parser.parse_args() + + fitsimage = FITSImage(args.infile) + fitsimage.crop(size=args.size) + fitsimage.write(args.outfile, clobber=args.clobber) + + +if __name__ == "__main__": + main() diff --git a/astro/fits/rescale_image.py b/astro/fits/rescale_image.py new file mode 100755 index 0000000..47db31f --- /dev/null +++ b/astro/fits/rescale_image.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2017 Aaron LI <aly@aaronly.me> +# MIT license +# + +""" +Rescale the FITS image to increase/decrease the pixel resolution +by interpolation, meanwhile keep the image FoV and update FITS WCS. + +For example, we may simulate the galaxy clusters (and point sources) +foreground component with a (much) higher resolution (smaller pixels) +compared to other foreground components (e.g., Galactic synchrotron) +and EoR signal. After simulating the observed images using OSKAR and +WSClean, the output images may in different sizes/resolution, which +may cause some troubles in subsequent usage, e.g., power spectrum +calculation. +""" + +import sys +import argparse +from datetime import datetime, timezone + +import numpy as np +from astropy.io import fits +from astropy.wcs import WCS +from scipy import ndimage + + +class FITSImage: + def __init__(self, filename): + self.filename = filename + self.header, self.image = self.open_image(filename) + print("Loaded FITS image from file: %s" % filename) + print("FITS image size: %dx%d" % (self.Nx, self.Ny)) + print("Pixel size: %.1f [arcsec]" % self.pixelsize) + print("FoV: %.1f [deg]" % self.fov[0]) + + def rescale(self, shape, order=1): + try: + Ny2, Nx2 = shape + except TypeError: + Ny2 = Nx2 = shape + print("Scale output size: %dx%d" % (Nx2, Ny2)) + print("Scale interpolation order: %d" % order) + zoom = ((Ny2+0.1)/self.Ny, (Nx2+0.1)/self.Nx) + pixelsize_old = self.pixelsize + self.image = ndimage.zoom(self.image, zoom=zoom, order=order) + self.pixelsize = pixelsize_old / zoom[0] + print("Scaled pixel size: %.1f [arcsec]" % self.pixelsize) + + def crop(self, size): + try: + xsize, ysize = size # [deg] + except TypeError: + xsize = ysize = size + Nx2 = xsize * 3600 / self.pixelsize + Ny2 = ysize * 3600 / self.pixelsize + if Nx2 > self.Nx or Ny2 > self.Ny: + raise ValueError("Crop region too large!") + + print("Central crop box size: %dx%d [deg]" % (xsize, ysize)) + print("Cropped image size: %dx%d" % (Nx2, Ny2)) + xi0 = int((self.Nx-Nx2) / 2) + yi0 = int((self.Ny-Ny2) / 2) + self.image = self.image[yi0:(yi0+Ny2), xi0:(xi0+Nx2)] + + def write(self, outfile, clobber=False): + header = self.header.copy(strip=True) + header.extend(self.wcs.to_header(), update=True) + header["DATE"] = (datetime.now(timezone.utc).astimezone().isoformat(), + "File creation date") + header.add_history(" ".join(sys.argv)) + hdu = fits.PrimaryHDU(data=self.data, header=header) + try: + hdu.writeto(outfile, overwrite=clobber) + except TypeError: + hdu.writeto(outfile, clobber=clobber) + print("Wrote scaled FITS image to file: %s" % outfile) + + @property + def fov(self): + # Unit: [deg] + return (self.Nx*self.pixelsize/3600, self.Ny*self.pixelsize/3600) + + @property + def pixelsize(self): + # Unit: [arcsec] + if hasattr(self, "_pixelsize"): + return self._pixelsize + else: + return self.header["CDELT1"] * 3600 # [deg] -> [arcsec] + + @pixelsize.setter + def pixelsize(self, value): + # Unit: [arcsec] + self._pixelsize = value + + @property + def Nx(self): + return self.image.shape[1] + + @property + def Ny(self): + return self.image.shape[0] + + @property + def wcs(self): + hdr = self.header + w = WCS(naxis=2) + w.wcs.equinox = hdr.get("EQUINOX", 2000.0) + w.wcs.ctype = [hdr.get("CTYPE1", "RA---SIN"), + hdr.get("CTYPE2", "DEC--SIN")] + w.wcs.crval = np.array([hdr.get("CRVAL1", 0.0), + hdr.get("CRVAL2", 0.0)]) + w.wcs.crpix = np.array([self.Ny/2+1, self.Nx/2+1]) + w.wcs.cdelt = np.array([self.pixelsize/3600, self.pixelsize/3600]) + w.wcs.cunit = [hdr.get("CUNIT1", "deg"), hdr.get("CUNIT2", "deg")] + return w + + @staticmethod + def open_image(infile): + """ + Open the slice image and return its header and 2D image data. + + NOTE + ---- + The input slice image may have following dimensions: + * NAXIS=2: [Y, X] + * NAXIS=3: [FREQ=1, Y, X] + * NAXIS=4: [STOKES=1, FREQ=1, Y, X] + + NOTE + ---- + Only open slice image that has only ONE frequency and ONE Stokes + parameter. + + Returns + ------- + header : `~astropy.io.fits.Header` + image : 2D `~numpy.ndarray` + The 2D [Y, X] image part of the slice image. + """ + with fits.open(infile) as f: + header = f[0].header + data = f[0].data + if data.ndim == 2: + # NAXIS=2: [Y, X] + image = data + elif data.ndim == 3 and data.shape[0] == 1: + # NAXIS=3: [FREQ=1, Y, X] + image = data[0, :, :] + elif data.ndim == 4 and data.shape[0] == 1 and data.shape[1] == 1: + # NAXIS=4: [STOKES=1, FREQ=1, Y, X] + image = data[0, 0, :, :] + else: + raise ValueError("Slice '{0}' has invalid dimensions: {1}".format( + infile, data.shape)) + return (header, image) + + +def main(): + parser = argparse.ArgumentParser( + description="Rescale a FITS image to the desired size/resolution") + parser.add_argument("-C", "--clobber", action="store_true", + help="overwrite existing output file") + parser.add_argument("--order", type=int, default=1, + help="scale interpolation order (default: 1)") + parser.add_argument("-s", "--size", type=int, required=True, + help="output image size (number of pixels)") + parser.add_argument("-i", "--infile", required=True, + help="input FITS image") + parser.add_argument("-o", "--outfile", required=True, + help="output FITS image") + args = parser.parse_args() + + fitsimage = FITSImage(args.infile) + fitsimage.rescale(shape=args.size, order=args.order) + fitsimage.write(args.outfile, clobber=args.clobber) + + +if __name__ == "__main__": + main() |