startracker.py

"""
startracker.py from UBNanoSat lib as modified by OpenLunar.

sg_2019
"""
from __future__ import print_function
from time import time
import sys
import traceback
import socket
import select
import os
import gc
import cv2
import numpy as np
from io import StringIO, BytesIO
import fcntl
from beast import beast
from systemd import daemon
import argparse

P_MATCH_THRESH = 0.99
SIMULATE = 0
if 'WATCHDOG_USEC' not in os.environ:
    os.environ['WATCHDOG_USEC'] = "30000000"


def trace(frame, event, arg):
    """Print a trace of some logging event."""
    print("%s, %s:%d" %
          (event, frame.f_code.co_filename, frame.f_lineno), file=sys.stderr)
    return trace


def setup(CONFIGFILE, YEAR):
    """Set up server before we do anything else."""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server.bind(('127.0.0.1', 8010))
    except:
        print("server socket already open: try terminal command: "
              "sudo kill $(sudo lsof -t -i:8010)")
        exit()

    server.listen(5)
    server.setblocking(0)

    print("Loading config")
    print (CONFIGFILE)
    beast.load_config(CONFIGFILE)
    print ("Loading hip_main.dat")
    S_DB = beast.star_db()
    S_DB.load_catalog("hip_main.dat", YEAR)
    print ("Filtering stars")
    SQ_RESULTS = beast.star_query(S_DB)
    SQ_RESULTS.kdmask_filter_catalog()
    SQ_RESULTS.kdmask_uniform_density(beast.cvar.REQUIRED_STARS)
    S_FILTERED = SQ_RESULTS.from_kdmask()
    print ("Generating DB")
    C_DB = beast.constellation_db(S_FILTERED, 2 + beast.cvar.DB_REDUNDANCY, 0)
    print ("Ready")
    return server


def a2q(att):
    """Attitude (3x3 matrix) converted to a quaternion."""
    q4 = 0.5 * np.sqrt(1 + np.trace(A))
    q1 = 1 / (4 * q4) * (att[1, 2] - att[2, 1])
    q2 = 1 / (4 * q4) * (att[2, 0] - att[0, 2])
    q3 = 1 / (4 * q4) * (att[0, 1] - att[1, 0])

    return np.array([q1, q2, q3, q4])


def q2a(q):
    """Quarternion converted to an attitude array."""
    q = q / np.linalg.norm(q)
    retval = np.array(
        [[
            q[0] ** 2 - q[1] ** 2 - q[2] ** 2 + q[3] ** 2,
            2 * (q[0] * q[1] + q[2] * q[3]),
            2 * (q[0] * q[2] - q[1] * q[3])],
            [
                2 * (q[0] * q[1] - q[2] * q[3]),
                - q[0] ** 2 + q[1] ** 2 - q[2] ** 2 + q[3] ** 2,
                2 * (q[1] * q[2] + q[0] * q[3])],
            [
                2 * (q[0] * q[2] + q[1] * q[3]),
                2 * (q[1] * q[2] - q[0] * q[3]),
                - q[0] ** 2 - q[1] ** 2 + q[2] ** 2 + q[3] ** 2]])
    return retval


def extrapolate_matrix(A, B, t1, t2, t3):
    """Calculate error angles between A and B via small angle approximation.

    of MRPs.
    """
    r = np.dot(B, np.transpose(A))
    dq = a2q(r)
    dp = np.array(
        [dq[0],
         dq[1],
         dq[2]]) / (1 + dq[3])
    angles_ab = 4 * dp

    # Extrapolate to new error angles between B and C.
    angles_bc = angles_ab / (t2 - t1) * (t3 - t2)

    # Convert to a quaternion via small angle approximation, then get C.
    c = q2a(np.array(
        [0.5 * angles_bc[0],
         0.5 * angles_bc[1],
         0.5 * angles_bc[2],
         1]))

    return (c, (1000000.0) * angles_ab / (t2 - t1))


# Note: SWIG's policy is to garbage collect objects created with
# constructors, but not objects created by returning from a function


def wahba(A, B, weight=[]):
    """
    Take in two matrices of points and finds the attitude matrix needed to.

    transform one onto the other.

    Input:
        A: nx3 matrix - x,y,z in body frame
        B: nx3 matrix - x,y,z in eci
        Note: the "n" dimension of both matrices must match
    Output:
        attitude_matrix: returned as a numpy matrix
    """
    assert len(A) == len(B)
    if (len(weight) == 0):
        weight = np.array([1] * len(A))
    # dot is matrix multiplication for array
    H = np.dot(np.transpose(A) * weight, B)

    # calculate attitude matrix
    # from http://malcolmdshuster.com/FC_MarkleyMortari_Girdwood_1999_AAS.pdf
    U, S, Vt = np.linalg.svd(H)
    flip = np.linalg.det(U) * np.linalg.det(Vt)

    U[:, 2] *= flip

    body2eci = np.dot(U, Vt)
    return body2eci


def print_ori(body2eci):
    """Print the orientation angles."""
    # rotation about the y axis (-90 to +90)
    dec = np.degrees(np.arcsin(body2eci[0, 2]))
    print ("DEC=" + str(dec), file=sys.stderr)

    # rotation about the z axis (-180 to +180)
    rotang = np.degrees(np.arctan2(body2eci[0, 1], body2eci[0, 0]))
    print ("RA=" + str(rotang), file=sys.stderr)

    # rotation about the camera axis (-180 to +180)
    orientation = np.degrees(-np.arctan2(body2eci[1, 2], body2eci[2, 2]))
    if orientation > 180:
        orientation = orientation - 360
    print ("ORIENTATION=" + str(orientation), file=sys.stderr)


class star_image:
    def __init__(self, imagefile, median_image):
        b_conf = [time(), beast.cvar.PIXSCALE, beast.cvar.BASE_FLUX]
        self.img_stars = beast.star_db()
        self.img_data = []
        self.match = None
        self.db_stars = None
        self.match_from_lm = None
        self.db_stars_from_lm = None

        # Placeholders so that these don't get garbage collected by SWIG
        self.fov_db = None
        self.const_from_lm = None

        # TODO: improve memory efficiency?? SG012019 BY WHOM?
        if "://" in imagefile:
            import urllib
            img = cv2.imdecode(
                np.asarray(
                    bytearray(
                        urllib.urlopen(imagefile).read()),
                    dtype="uint8"),
                cv2.IMREAD_COLOR)
        else:
            img = cv2.imread(imagefile)
        if img is None:
            print ("Invalid image, using blank dummy image", file=sys.stderr)
            img = median_image

        img = np.clip(
            img.astype(np.int16) - median_image,
            a_min=0,
            a_max=255).astype(np.uint8)
        img_grey = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

        # removes areas of the image that don't meet our brightness threshold
        ret, thresh = cv2.threshold(
            img_grey,
            beast.cvar.THRESH_FACTOR * beast.cvar.IMAGE_VARIANCE,
            255,
            cv2.THRESH_BINARY)
        contours, heirachy = cv2.findContours(thresh, 1, 2)

        for c in contours:
            M = cv2.moments(c)
            if M['m00'] > 0:
                # this is how the x and y position are defined by cv2
                cx = M['m10'] / M['m00']
                cy = M['m01'] / M['m00']
                # see https://alyssaq.github.io/2015/
                # computing-the-axes-or-orientation-of-a-blob/
                # for how to convert these into eigenvectors/values
                u20 = M["m20"] / M["m00"] - cx**2
                u02 = M["m02"] / M["m00"] - cy**2
                u11 = M["m11"] / M["m00"] - cx * cy
                # the center pixel is used as the approximation of
                # the brightest pixel
                self.img_stars += beast.star(cx - beast.cvar.IMG_X / 2.0,
                                             (cy - beast.cvar.IMG_Y / 2.0),
                                             float(cv2.getRectSubPix(img_grey,
                                                   (1, 1),
                                                   (cx, cy))[0, 0]),
                                             -1)
                self.img_data.append(b_conf +
                                     [cx, cy, u20, u02, u11] +
                                     cv2.getRectSubPix(img,
                                                       (1, 1),
                                                       (cx, cy)
                                                       )[0, 0].tolist())

    def match_near(self, x, y, z, r):
        SQ_RESULTS.kdsearch(x, y, z, r,
                            beast.cvar.THRESH_FACTOR *
                            beast.cvar.IMAGE_VARIANCE)
        # estimate density for constellation generation
        C_DB.results.kdsearch(x, y, z, r,
                              beast.cvar.THRESH_FACTOR *
                              beast.cvar.IMAGE_VARIANCE)
        fov_stars = SQ_RESULTS.from_kdresults()
        self.fov_db = beast.constellation_db(fov_stars,
                                             C_DB.results.r_size(),
                                             1)
        C_DB.results.clear_kdresults()
        SQ_RESULTS.clear_kdresults()

        img_const = beast.constellation_db(self.img_stars,
                                           beast.cvar.MAX_FALSE_STARS + 2,
                                           1)
        near = beast.db_match(self.fov_db, img_const)
        if near.p_match > P_MATCH_THRESH:
            self.match = near
            self.db_stars = near.winner.from_match()

    def match_lis(self):
        """For the first pass, we only want to use the brightest.

         MAX_FALSE_STARS+REQUIRED_STARS."""
        img_stars_n_brightest = \
            self.img_stars.copy_n_brightest(beast.cvar.MAX_FALSE_STARS +
                                            beast.cvar.REQUIRED_STARS)
        img_const_n_brightest = \
            beast.constellation_db(img_stars_n_brightest,
                                   beast.cvar.MAX_FALSE_STARS + 2,
                                   1)
        lis = beast.db_match(C_DB, img_const_n_brightest)
        # TODO: uncomment once p_match is fixed
        # if lis.p_match>P_MATCH_THRESH:
        if (lis.p_match > P_MATCH_THRESH and
                lis.winner.size() >= beast.cvar.REQUIRED_STARS):
            x = lis.winner.R11
            y = lis.winner.R21
            z = lis.winner.R31
            self.match_near(x, y, z, beast.cvar.MAXFOV / 2)

    def match_rel(self, last_match):
        """Match stars relative to something?."""
        # make copy of stars from lastmatch
        img_stars_from_lm = last_match.img_stars.copy()
        w = last_match.match.winner
        # convert the stars to ECI
        for i in range(img_stars_from_lm.size()):
            s = img_stars_from_lm.get_star(i)
            x = s.x * w.R11 + s.y * w.R12 + s.z * w.R13
            y = s.x * w.R21 + s.y * w.R22 + s.z * w.R23
            z = s.x * w.R31 + s.y * w.R32 + s.z * w.R33
            s.x = x
            s.y = y
            s.z = z
        # create constellation from last match
        self.const_from_lm = \
            beast.constellation_db(img_stars_from_lm,
                                   beast.cvar.MAX_FALSE_STARS + 2,
                                   1)
        # match between last and current
        img_const = beast.constellation_db(self.img_stars,
                                           beast.cvar.MAX_FALSE_STARS + 2,
                                           1)
        rel = beast.db_match(self.const_from_lm, img_const)
        if rel.p_match > P_MATCH_THRESH:
            self.match_from_lm = rel
            self.db_stars_from_lm = rel.winner.from_match()

    def print_match(self, body_correction=None, angrate_string=""):
        """Print stellar match for debug."""
        if body_correction is None:
            body_correction = np.eye(3)
        if self.match is not None:
            self.match.winner.print_ori()
        db = self.db_stars
        im = self.img_stars
        if db is None:
            if self.db_stars_from_lm is None:
                # neither relative nor absolute matching could be used
                print("")
                return
            else:
                db = self.db_stars_from_lm
        assert(db.size() == im.size())
        star_out = []
        for i in range(db.size()):
            s_im = im.get_star(i)
            s_db = db.get_star(i)
            if (s_db.id >= 0):
                weight = 1.0 / (s_db.sigma_sq + s_im.sigma_sq)
                temp = np.dot(bodyCorrection,
                              np.array([[s_im.x], [s_im.y], [s_im.z]]))
                star_out.append(str(temp[0, 0]) + ',' +
                                str(temp[1, 0]) + ',' +
                                str(temp[2, 0]) + ',' +
                                str(s_db.x) + ',' +
                                str(s_db.y) + ',' +
                                str(s_db.z) + ',' +
                                str(weight))
        print("stars", len(star_out), file=sys.stderr)
        print("ang_rate: " + angrate_string, file=sys.stderr)
        print(" ".join(star_out) + " " + angrate_string)

NONSTARS = {}
NONSTAR_NEXT_ID = 0
NONSTAR_DATAFILENAME = "/dev/null"
# NONSTAR_DATAFILENAME="data"+str(time())+".txt"
NONSTAR_DATAFILE = open(NONSTAR_DATAFILENAME, "w")


class nonstar:
    def __init__(self, current_image, i, source):
        global NONSTARS, NONSTAR_NEXT_ID, NONSTAR_DATAFILENAME, NONSTAR_DATAFILE
        self.id=NONSTAR_NEXT_ID
        NONSTARS[self.id] = self
        current_image.img_stars.get_star(i).id = self.id
        NONSTAR_NEXT_ID += 1
        self.data = []
        self.add_data(current_image, i, source)

    def add_data(self, current_image, i, source):
        s_im = current_image.img_stars.get_star(i)
        s_db_x = 0.0
        s_db_y = 0.0
        s_db_z = 0.0
        w = None
        if (current_image.match is not None and
                current_image.match.p_match > P_MATCH_THRESH):
            w = current_image.match.winner
        elif (current_image.match is not None and
                current_image.match.p_match > P_MATCH_THRESH):
            w = current_image.match_from_lm.winner
        if w is not None:
            # convert the stars to ECI
            s_db_x = s_im.x * w.R11 + s_im.y * w.R12 + s_im.z * w.R13
            s_db_y = s_im.x * w.R21 + s_im.y * w.R22 + s_im.z * w.R23
            s_db_z = s_im.x * w.R31 + s_im.y * w.R32 + s_im.z * w.R33
        self.data.append([source,
                          s_im.x, s_im.y, s_im.z,
                          s_db_x, s_db_y, s_db_z] +
                         current_image.img_data[i])

    def write_data(self, fd):
        if sys.version_info[0] > 2:
            os.write(fd, bytes(str(self.id) + " " +
                               str(len(self.data)) + "\n",
                               encoding='UTF-8'))
        else:
            os.write(fd, str(self.id) + " " + str(len(self.data)) + "\n")
        for i in self.data:
            s = [str(j) for j in i]
        if sys.version_info[0] > 2:
            os.write(fd, bytes(" ".join(s) + "\n", encoding='UTF-8'))
        else:
            os.write(fd, " ".join(s) + "\n")

    def __del__(self):
        self.write_data(NONSTAR_DATAFILE.fileno())


def flush_nonstars():
    """Clear global variables, close files, call gc."""
    global NONSTARS, NONSTAR_NEXT_ID, NONSTAR_DATAFILENAME, NONSTAR_DATAFILE
    NONSTARS = {}
    NONSTAR_NEXT_ID = 0
    gc.collect()
    NONSTAR_DATAFILE.close()
    NONSTAR_DATAFILENAME = "data" + str(time()) + ".txt"
    NONSTAR_DATAFILE = open(NONSTAR_DATAFILENAME, "w")


def update_nonstars(current_image, source):
    """Iterate over all stars, tag nonstars."""
    global NONSTARS, NONSTAR_NEXT_ID
    nonstars_next = {}
    im = current_image.img_stars
    db = current_image.db_stars
    db_lm = current_image.db_stars_from_lm
    if db is not None:
        assert(db.size() == im.size())
    if db_lm is not None:
        assert(db_lm.size() == im.size())
        for i in range(im.size()):
            im.get_star(i).id = db_lm.get_star(i).id
    for i in range(im.size()):
        s_im = im.get_star(i)
        # is this a star? if so remove from nonstars
        if db is not None and db.get_star(i).id >= 0:
            if (s_im.id in NONSTARS):
                del NONSTARS[s_im.id]
            s_im.id = -1
        # if it's already there, add the latest mesurement
        elif (s_im.id in NONSTARS):
            NONSTARS[s_im.id].add_data(current_image, i, source)
            nonstars_next[s_im.id] = NONSTARS[s_im.id]
        # otherwise add a new nonstar
        else:
            ns = nonstar(current_image, i, source)
            nonstars_next[ns.id] = ns
    NONSTARS = nonstars_next

    # wrap around to prevent integer overflow
    if (NONSTAR_NEXT_ID > 2**30):
        flush_nonstars()


def winner_attitude(w):
    """Find attitude of winning match."""
    # w=self.last_match.match.winner
    eci2body = np.array(
        [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]],
        dtype=float)
    eci2body[0, 0] = w.R11
    eci2body[0, 1] = w.R12
    eci2body[0, 2] = w.R13
    eci2body[1, 0] = w.R21
    eci2body[1, 1] = w.R22
    eci2body[1, 2] = w.R23
    eci2body[2, 0] = w.R31
    eci2body[2, 1] = w.R32
    eci2body[2, 2] = w.R33
    return np.transpose(eci2body)


class StarCamera:
    """Definition for camera hardware and image source."""

    def __init__(self, median_file, source="RGB"):
        self.source = source
        self.current_image = None
        self.last_match = None
        self.median_image = cv2.imread(median_file)

    def solve_image(self, imagefile, lis=1, quiet=0):
        starttime = time()
        if SIMULATE == 1 and quiet == 0:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.connect(("jeb", 7011))
            data = s.recv(2048)
            s.close()
        print("Time1: " + str(time() - starttime), file=sys.stderr)
        self.current_image = star_image(imagefile, self.median_image)
        print("Time2: " + str(time() - starttime), file=sys.stderr)
        if (lis == 1):
            self.current_image.match_lis()
        print("Time3: " + str(time() - starttime), file=sys.stderr)
        if self.last_match is not None:
            self.current_image.match_rel(self.last_match)
        if quiet == 0:
            if SIMULATE == 1:
                print (data.rstrip("\n").rstrip("\r"))
            else:
                self.current_image.print_match()
            print("Time4: " + str(time() - starttime), file=sys.stderr)

        update_nonstars(self.current_image, self.source)
        print("Time5: " + str(time() - starttime), file=sys.stderr)
        if self.current_image.match is not None:
            self.last_match = self.current_image
        else:
            self.last_match = None
        print("Time6: " + str(time() - starttime), file=sys.stderr)

    def extrapolate_image(self, imagefile1, imagefile2, time1, time2):
        # self.solve_image(imagefile2,lis=1,quiet=0)
        self.solve_image(imagefile1, lis=1, quiet=1)
        print(1, file=sys.stderr)
        if (self.last_match is None):
            print(2, file=sys.stderr)
            print ("")
            return
        a1 = winner_attitude(self.last_match.match.winner)
        self.solve_image(imagefile2, lis=1, quiet=1)
        print(3, file=sys.stderr)
        if (self.last_match is None):
            print(4, file=sys.stderr)
            print ("")
            return
        a2 = winner_attitude(self.last_match.match.winner)
        print(a1, a2, np.linalg.svd(a1)[1],
              np.linalg.svd(a1)[1], file=sys.stderr)
        a, angrate = extrapolate_matrix(a1, a2, time1, time2, time() * 1e6)
        print(a, np.linalg.svd(a)[1], file=sys.stderr)

        self.last_match.print_match(
            a, ",".join([str(i) for i in angrate.tolist()]))


# dummy for now
# TODO: add science data from IR cam
class science_camera:
    def __init__(self, median_file, source="IR"):
        self.source = source
        self.current_image = None
        self.last_match = None
        self.median_image = cv2.imread(median_file)

    def solve_image(self, imagefile):
        if sys.version_info[0] > 2:
            os.write(1, bytes(os.path.abspath(NONSTAR_DATAFILENAME),
                     encoding='UTF-8'))
        else:
            os.write(1, os.path.abspath(NONSTAR_DATAFILENAME))




class connection:
    """Tracks activity on a file descriptor and allows TCP read/writes."""
    def __init__(self, conn, epoll):
        """
        Create connection to track file descriptor activity.

        @note: Adds C{fd . self} to C{CONNECTIONS}
        @param conn: Any file object with the fileno() method
        @param epoll: File descriptor edge polling object
        """
        self.conn = conn
        self.fd = self.conn.fileno()
        epoll.register(self.fd, select.EPOLLIN)
        self.epoll = epoll
        CONNECTIONS[self.fd] = self

    def read(self):
        """
        Complete non-blocking read on file descriptor.

        of an arbitrary amount of data
        @return: Entire read string
        @rtype: C{string}
        """
        # need nonblocking for read
        fl = fcntl.fcntl(self.fd, fcntl.F_GETFL)
        fcntl.fcntl(self.fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        data = b''
        try:
            while True:
                lastlen = len(data)
                data += os.read(self.fd, 1024)
                if len(data) == lastlen:
                    break
        except OSError as e:
            # error 11 means we have no more data to read
            if e.errno == 11:
                pass
            elif e.errno == 104:
                print("WARNING: ABNORMAL DISCONNECT", file=sys.stderr)
            else:
                raise
        return data

    def write(self, data):
        """
        Blocking read on file descriptor.

        @param data: ASCII data to write
        @type data: C{string}
        """
        if len(data) == 0:
            return
        if self.fd == 0:
            if sys.version_info[0] > 2:
                os.write(1, bytes(data, encoding='UTF-8'))
            else:
                os.write(1, data)
            return
        # need blocking IO for writing
        fl = fcntl.fcntl(self.fd, fcntl.F_GETFL)
        fcntl.fcntl(self.fd, fcntl.F_SETFL, fl & ~os.O_NONBLOCK)
        if sys.version_info[0] > 2:
            os.write(self.fd, bytes(data, encoding='UTF-8'))
        else:
            os.write(self.fd, data)

    def close(self):
        """Close connection with file descriptor."""
        self.epoll.unregister(self.fd)
        self.conn.close()
        if CONNECTIONS[self.fd]:
            del CONNECTIONS[self.fd]


def parse_args():
    """Parse args.

    from
    CONFIGFILE = sys.argv[1]
    YEAR = float(sys.argv[2])
    """
    parser = argparse.ArgumentParser(description="""
           Starcam solver entrypoint.

        """)
    parser.add_argument('--configfile', action="store", dest="CONFIGFILE")
    parser.add_argument('--year', action="store", dest="YEAR", type=int)
    parser.add_argument('--timeout', action="store", dest="WATCHDOG_USEC", type=int)
    parser.add_argument('--camera', action="store", dest="CAM")
    return parser.parse_args()


def main(args):
    """It is a main.

    arguments are passed in via env vars currently;
        WATCHDOG_USEC = 3000000
    """
    server = setup(args.CONFIGFILE, args.YEAR)
    rgb = StarCamera(args.CAM)
    ir = science_camera(args.CAM)

    CONNECTIONS = {}

    epoll = select.epoll()
    epoll.register(server.fileno(), select.EPOLLIN)

    try:
        connection(sys.stdin, epoll)
    except IOError:
        pass

    daemon.notify("WATCHDOG=1")
    last_ping = time()
    while True:
        # systemd watchdog
        events = epoll.poll(float(os.environ['WATCHDOG_USEC']) /
                            2.0e6 - (time() - last_ping))
        if (len(events) == 0 or time() >=
                (last_ping + float(os.environ['WATCHDOG_USEC']) / 2.0e6)):
            daemon.notify("WATCHDOG=1")
            last_ping = time()
        # events = epoll.poll()
        for fd, event_type in events:
            # Activity on the master socket means a new connection.
            if fd == server.fileno():
                conn, addr = server.accept()
                connection(conn, epoll)
            elif fd in CONNECTIONS:
                w = CONNECTIONS[fd]
                data = w.read()
                print(data.decode(encoding='UTF-8'), file=sys.stderr)
                if len(data) > 0:
                    if sys.version_info[0] > 2:
                        stdout_redir = StringIO()
                    else:
                        stdout_redir = BytesIO()
                    stdout_old = sys.stdout
                    sys.stdout = stdout_redir
                    try:
                        exec(data)
                    except SystemExit:
                        w.close()
                        raise
                    except:
                        traceback.print_exc(file=sys.stdout)
                    sys.stdout = stdout_old
                    data_out = stdout_redir.getvalue()
                    print(data_out, file=sys.stderr)
                    w.write(data_out)
                else:
                    w.close()


if __name__ == "__main__":
    args = parse_args()
    main(args)