From 11927ca2d5481d4abe2d3680897e984af5e65e68 Mon Sep 17 00:00:00 2001
From: Wilfred Tyler Gee <wtylergee@gmail.com>
Date: Fri, 29 Dec 2017 23:58:48 +1100
Subject: [PATCH] Move polar alignment helper functions into POCS images utils
 (#265)

---
 bin/pocs_shell       |   4 +-
 pocs/utils/images.py | 108 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 106 insertions(+), 6 deletions(-)

diff --git a/bin/pocs_shell b/bin/pocs_shell
index 43b557bd0..b96ae2552 100755
--- a/bin/pocs_shell
+++ b/bin/pocs_shell
@@ -26,8 +26,6 @@ from pocs.utils import listify
 from pocs.utils.database import PanMongo
 from pocs.utils.messaging import PanMessaging
 
-from piaa.utils import polar_alignment
-
 
 class PocsShell(Cmd):
 
@@ -153,7 +151,7 @@ class PocsShell(Cmd):
 
 simulate is a space-separated list of hardware to simulate.
 Hardware names: {}   (or all for all hardware)'''.format(
-','.join(hardware.get_all_names())))
+            ','.join(hardware.get_all_names())))
 
     def do_reset_pocs(self, *arg):
         self.pocs = None
diff --git a/pocs/utils/images.py b/pocs/utils/images.py
index 1d26e34a0..3e7505637 100644
--- a/pocs/utils/images.py
+++ b/pocs/utils/images.py
@@ -12,12 +12,19 @@
 
 import numpy as np
 
-from ffmpy import FFmpeg
-from glob import glob
+from skimage.feature import canny
+from skimage.transform import hough_circle
+from skimage.transform import hough_circle_peaks
 
-from astropy import units as u
 from astropy.io import fits
+from astropy.nddata import Cutout2D
+from astropy.visualization import SqrtStretch
+from astropy.visualization.mpl_normalize import ImageNormalize
 from astropy.wcs import WCS
+from astropy import units as u
+
+from ffmpy import FFmpeg
+from glob import glob
 
 from pocs.utils import current_time
 from pocs.utils import error
@@ -1019,3 +1026,98 @@ def clean_observation_dir(dir_name, *args, **kwargs):
     except Exception as e:
         warn(
             'Problem with cleanup creating timelapse:'.format(e))
+
+
+def analyze_polar_rotation(pole_fn, *args, **kwargs):
+    """ Get celestial pole XY coordinates
+
+    Args:
+        pole_fn (str): FITS file of celestial pole
+
+    Returns:
+        tuple(int): A tuple of integers corresponding to the XY pixel position
+        of celestial pole
+    """
+    get_solve_field(pole_fn, **kwargs)
+
+    wcs = WCS(pole_fn)
+
+    pole_cx, pole_cy = wcs.all_world2pix(360, 90, 1)
+
+    return pole_cx, pole_cy
+
+
+def analyze_ra_rotation(rotate_fn):
+    """ Get RA axis center of rotation XY coordinates
+
+    Args:
+        rotate_fn (str): FITS file of RA rotation
+
+    Returns:
+        tuple(int): A tuple of integers corresponding to the XY pixel position
+        of the center of rotation around RA axis
+    """
+    d0 = fits.getdata(rotate_fn)  # - 2048
+
+    # Get center
+    position = (d0.shape[1] // 2, d0.shape[0] // 2)
+    size = (1500, 1500)
+    d1 = Cutout2D(d0, position, size)
+
+    d1.data = d1.data / d1.data.max()
+
+    # Get edges for rotation
+    rotate_edges = canny(d1.data, sigma=1.0)
+
+    rotate_hough_radii = np.arange(100, 500, 50)
+    rotate_hough_res = hough_circle(rotate_edges, rotate_hough_radii)
+    rotate_accums, rotate_cx, rotate_cy, rotate_radii = \
+        hough_circle_peaks(rotate_hough_res, rotate_hough_radii, total_num_peaks=1)
+
+    return d1.to_original_position((rotate_cx[-1], rotate_cy[-1]))
+
+
+def plot_center(pole_fn, rotate_fn, pole_center, rotate_center):
+    """ Overlay the celestial pole and RA rotation axis images
+
+    Args:
+        pole_fn (str): FITS file of polar center
+        rotate_fn (str): FITS file of RA rotation image
+        pole_center (tuple(int)): Polar center XY coordinates
+        rotate_center (tuple(int)): RA axis center of rotation XY coordinates
+
+    Returns:
+        matplotlib.Figure: Plotted image
+    """
+    d0 = fits.getdata(pole_fn) - 2048.  # Easy cast to float
+    d1 = fits.getdata(rotate_fn) - 2048.  # Easy cast to float
+
+    d0 /= d0.max()
+    d1 /= d1.max()
+
+    pole_cx, pole_cy = pole_center
+    rotate_cx, rotate_cy = rotate_center
+
+    d_x = pole_center[0] - rotate_center[0]
+    d_y = pole_center[1] - rotate_center[1]
+
+    fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(20, 14))
+
+    # Show rotation center in red
+    ax.scatter(rotate_cx, rotate_cy, color='r', marker='x', lw=5)
+
+    # Show polar center in green
+    ax.scatter(pole_cx, pole_cy, color='g', marker='x', lw=5)
+
+    # Show both images in background
+    norm = ImageNormalize(stretch=SqrtStretch())
+    ax.imshow(d0 + d1, cmap='Greys_r', norm=norm, origin='lower')
+
+    # Show an arrow
+    if (np.abs(pole_cy - rotate_cy) > 25) or (np.abs(pole_cx - rotate_cx) > 25):
+        ax.arrow(rotate_cx, rotate_cy, pole_cx - rotate_cx, pole_cy -
+                 rotate_cy, fc='r', ec='r', width=20, length_includes_head=True)
+
+    ax.set_title("dx: {:0.2f} pix      dy: {:0.2f} pix".format(d_x, d_y))
+
+    return fig