Add cutout image

1 files changed, 34 insertions, 14 deletions

diff --git a/get_image.py b/get_image.py
index 83577a8..3f7cb8d 100644
--- a/get_image.py
+++ b/get_image.py
@@ -20,34 +20,54 @@ Excample:
 import sys
 
 import astropy.units as u
+from astropy.coordinates import SkyCoord
 from astroquery.sdss import SDSS
+from astropy.wcs import WCS
+import matplotlib as mpl
+import matplotlib.pyplot as plt
 
 
 def main():
 
     # Init
-    ra = sys.argv[1]
-    dec = sys.argv[2]
+    ra = float(sys.argv[1])
+    dec = float(sys.argv[2])
     band = sys.argv[3]
 
     # Query the region
-    result = SDSS.query_region(ra + " " + dec,
+    coord = SkyCoord(ra=ra*u.deg, dec=dec*u.deg)
+    result = SDSS.query_region(coord,
                                radius=1 * u.arcmin,
                                data_release=12)
 
     # Fetch the first image of the query
     result = result[0]
-    im = SDSS.get_images(run=result["run"],
-                         camcol=result["camcol"],
-                         field=result["field"],
-                         data_release=12,
-                         band=band)
-
-    # Plot the first image as an example
-    import matplotlib as mpl
-    import matplotlib.pyplot as plt
-    im = im[0][0].data
-    plt.imshow(im, origin="lower")
+    image = SDSS.get_images(run=result["run"],
+                            camcol=result["camcol"],
+                            field=result["field"],
+                            data_release=12,
+                            band=band)
+
+    # Read in data, header, and a dx x dy sized cutout image. Note that the WCS
+    # solution will only work on the original image.
+    data = image[0][0].data
+    header = image[0][0].header
+    wcs = WCS(header)
+    x, y = map(int, wcs.world_to_pixel(coord))
+    dx = dy = 100
+    cutout = data[y - dy // 2 : y + dy // 2,
+                  x - dx // 2 : x + dx // 2]
+
+    # Plot the first image, and a cutout of it as an example
+    plt.figure(1)
+    plt.imshow(data, origin="lower", norm=mpl.colors.LogNorm())
+    plt.plot(x, y, 'rx')
+    plt.gca().add_patch(
+        plt.Rectangle([x - dx / 2, y - dy / 2], dx, dy, color="red", fc="none")
+    )
+
+    plt.figure(2)
+    plt.imshow(cutout, norm=mpl.colors.LogNorm(), origin="lower")
     plt.show()