Added tests and fictrac code

Basic tests for most of the code complete. .gitignore was blocking
fictrac folders from being tracked as well.

Author: davidt0x

.gitignore

@@ -3,16 +3,11 @@
 
 *.pyc
 
-nlopt-2.4.2-dll64/
-
 example/*.dat
 example/*.log
 example/test-template.png
 example/output.txt
 
-fictrac/
-
-
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

example/run_example.py

@@ -13,10 +13,9 @@ def run_pybmt_example():
     # Instantiate a FicTracDriver object to handle running of FicTrac in the background and communication
     # of program state.
     tracDrv = FicTracDriver(config_file=fictrac_config, console_ouput_file=fictrac_console_out,
-                            #remote_endpoint_url="localhost:5556",
                             track_change_callback=callback, plot_on=False)
 
-    # This will start FicTrac and the
+    # This will start FicTrac and block until complete
     tracDrv.run()
 
 if __name__ == "__main__":

pybmt/fictrac/__init__.py

@@ -0,0 +1,230 @@
+import subprocess
+import time
+import os
+from collections import deque
+
+import zmq
+
+from pybmt.fictrac.state import FicTracState
+from pybmt.tools import which
+
+
+class FicTracDriver:
+    """
+    This class drives the tracking of the fly via a separate software called FicTrac. It invokes this process and
+    calls a control function once for each time the tracking state is updated.
+    """
+    def __init__(self, config_file=None, remote_endpoint_url=None, console_ouput_file="output.txt",
+                 track_change_callback=None, pgr_enable=False, plot_on=True, fic_trac_bin_path=None):
+        """
+        Create the FicTrac driver object. This function will perform a check to see if the FicTrac program is present
+        on the path. If it is not, it will throw an exception.
+
+        :param str config_file: The path to the configuration file that should be passed to the FicTrac command.
+        :param str remote_endpoint_url: If FicTrac is already running on another machine, this is the url.
+        :param str console_output_file: The path to the file where console output should be stored.
+        :param track_change_callback: A FlyVRCallback class which is called once everytime tracking status changes. See
+        control.FlyVRCallback for example.
+        :param bool pgr_enable: Is Point Grey camera support needed. This just decides which executable to call, either
+        'FicTrac' or 'FicTrac-PGR'.
+        :param str fic_trac_bin_path: The path the the fictrac binary to use. Default is None. If None, we will try to
+        find fictrac on the path.
+        :param str remote_enpoint_url
+        """
+
+        self.track_change_callback = track_change_callback
+        self.plot_on = plot_on
+
+        # The message loop has to stay above this average number of frames per second. If it falls below we are
+        # not grabbing messages fast enough and will fall behind FicTrac in state. I don't like this solution that
+        # much, with shared memory this was easier to detect.
+        self.average_fps_threshold = 400
+
+        # If fictrac is already running, for example, on another machine, then we don't need to worry about running it.
+        if remote_endpoint_url is not None:
+            self.remote_endpoint_url = "tcp://" + remote_endpoint_url
+            self.start_fictrac = False
+        else:
+            self.start_fictrac = True
+            self.remote_endpoint_url = "tcp://localhost:5556"
+
+            self.config_file = config_file
+
+            # Get the directory that the config file is in, this will be the current working directory
+            # of FicTrac.
+            self.config_dir = os.path.dirname(self.config_file)
+            if self.config_dir == "":
+                self.config_dir = None
+
+            # Get base config file name
+            self.config_file_base = os.path.basename(self.config_file)
+
+            self.console_output_file = console_ouput_file
+            self.track_change_callback = track_change_callback
+            self.pgr_enable = pgr_enable
+            self.plot_on = plot_on
+
+            # If the user didn't specify the path to fictrac, look for it on the path.
+            if fic_trac_bin_path is None:
+                self.fictrac_bin = 'fictrac'
+                if self.pgr_enable:
+                    self.fictrac_bin = 'fictrac-pgr'
+
+                # If this is Windows, we need to add the .exe extension.
+                if os.name == 'nt':
+                    self.fictrac_bin = self.fictrac_bin + ".exe"
+
+                # Lets make sure FicTrac exists on the path
+                self.fictrac_bin_fullpath = which(self.fictrac_bin)
+
+                if self.fictrac_bin_fullpath is None:
+                    raise RuntimeError("Could not find " + self.fictrac_bin + " on the PATH!")
+
+            else:
+                self.fictrac_bin_fullpath = fic_trac_bin_path
+
+                # TODO: Make sure we are using the correct version of fictrac.
+
+        self.fictrac_process = None
+
+    def run(self):
+        """
+        Start the the FicTrac process and block till it closes. This function will poll a shared memory region for
+        changes in tracking data and invoke a control function when they occur. FicTrac is assumed to exist on the
+        system path.
+
+        :return:
+        """
+
+        # Setup anything the callback needs.
+        self.track_change_callback.setup_callback()
+
+        try:
+            # Start FicTrac if we need to
+            if self.start_fictrac:
+                with open(self.console_output_file, "wb") as out:
+                    self.fictrac_process = subprocess.Popen([self.fictrac_bin_fullpath, "--exit_when_done",
+                                                         self.config_file_base],
+                                                         stdout=out, stderr=subprocess.STDOUT,
+                                                         cwd=self.config_dir)
+                    self._process_messages()
+            else:
+                self._process_messages()
+
+            # Call poll one last time to get the return value
+            self.fictrac_process.poll()
+
+            # Get the fic trac process return code
+            if self.fictrac_process is not None and self.fictrac_process.returncode is not None and self.fictrac_process.returncode != 0:
+                raise RuntimeError("FicTrac failed because of an application error. " +
+                                   "Consult the FicTrac console output file ({}). ".format(self.console_output_file))
+            if self.frame_cnt == 0:
+                raise RuntimeError("Zero frames processed. FicTrac failed because of an application error. " +
+                             "Consult the FicTrac console output file ({}). ".format(self.console_output_file))
+
+
+        except Exception as ex:
+            if self.fictrac_process is not None:
+                self.fictrac_process.terminate()
+
+            raise Exception("PyBMT Error!") from ex
+        finally:
+            self.track_change_callback.shutdown_callback()
+            self._cleanup()
+
+    def _process_messages(self):
+
+        #  Setup ZeroMQ context and socket to talk to server
+        context = zmq.Context()
+        socket = context.socket(zmq.SUB)
+
+        # Set a timeout on this socket so we don't block forever
+        socket.RCVTIMEO = 1000  # in milliseconds
+
+        # This is the receiver high water mark, zero mq will start to drop incoming messages after it
+        # has queued this many. This will let us detect if we are not picking up messages quick enough because of a
+        # slow callback process. This isn't perfect though since OS buffers messages as well.
+        socket.setsockopt(zmq.RCVHWM, 1)
+
+        # Bind and subscribe
+        socket.connect(self.remote_endpoint_url)
+        socket.setsockopt(zmq.SUBSCRIBE, b"")
+
+        #           if self.plot_on:
+        #               self.plot_task = ConcurrentTask(task=plot_task_fictrac, comms="pipe",
+        #                                      taskinitargs=[state])
+        #               self.plot_task.start()
+
+        # Lets track average frame rate we are receiving the messages
+        # Our buffer of past speeds
+        time_history = deque(maxlen=10)
+        avg_fps = 0
+        self.frame_cnt = 0
+
+        # Lets keep track of the last fictrac state we received
+        last_fstate: FicTracState = None
+        fstate: FicTracState = None
+        isOK = True
+        while isOK:
+
+            # If we are running FicTrac, make sure it is running still
+            #if frame_cnt > 0 and self.fictrac_process is not None and self.fictrac_process.poll() is None:
+            #    print("FicTrac process gone!")
+            #    break
+
+            # Receive state update from FicTrac process
+            try:
+                data = socket.recv_string()
+            except zmq.error.Again:
+
+                # If we get socket error, probably means fictrac is gone.  If we started it, just break.
+                # If we didn't start it, signal the connection error.
+                if self.start_fictrac:
+                    break
+                else:
+                    raise Exception("Socket timed out. Couldn't reach fictrac!")
+
+            # Message received start the timer, want to keep track of how long it takes to process the message.
+            t0 = time.perf_counter()
+
+            # If FicTrac sent and END signal, its time to clean up
+            if data == "END":
+                break
+
+            # Lets keep track of the last fictrac state we received
+            last_fstate = fstate
+
+            # Parse the data packet into our state structure. Get our new state.
+            fstate = FicTracState.zmq_string_msg_to_state(data)
+
+            if last_fstate is not None and fstate.frame_cnt - last_fstate.frame_cnt != 1:
+                self.fictrac_process.terminate()
+                raise Exception(("FicTrac frame counter jumped by more than 1! oldFrame = " +
+                                 str(last_fstate.frame_cnt) + ", newFrame = " + str(fstate.frame_cnt)))
+
+            # Call the main callback function with the current state
+            isOK = self.track_change_callback.process_callback(fstate)
+
+            # Stop the clock
+            t1 = time.perf_counter()
+
+            # Add the speed to our history
+            time_history.append(t1 - t0)
+
+            # Get the running average fps
+            avg_fps = 1 / (sum(time_history) / len(time_history))
+
+            if avg_fps < self.average_fps_threshold and frame_cnt > 300:
+                self.fictrac_process.terminate()
+                raise Exception("Average FPS fell below avg_fps_threshold({}). Processing callback is " +
+                                "probably operating too slow.")
+
+            self.frame_cnt = self.frame_cnt + 1
+
+    def _cleanup(self):
+        """
+        This is method is called whenever PyBMT run is shutting down things.
+
+        :return:
+        """
+        pass

pybmt/fictrac/config.py

@@ -0,0 +1,112 @@
+
+
+def angle_diff(angle1, angle2):
+    diff = angle2 - angle1
+    while diff < np.deg2rad(-180.0):
+        diff += np.deg2rad(360.0)
+    while diff > np.deg2rad(180):
+        diff -= np.deg2rad(360.0)
+    return diff
+
+def plot_task_fictrac(remote_endpoint_url,
+                      fictrac_state_fields=['speed', 'direction', 'del_rot_cam_vec', 'del_rot_error'],
+                      num_history=200):
+    """
+    A coroutine for plotting fast, realtime as per: https://gist.github.com/pklaus/62e649be55681961f6c4. This is used
+    for plotting streaming data coming from FicTrac.
+
+    :param disp_queue: The message queue from which data is sent for plotting.
+    :return: None
+    """
+
+    # Open or create the shared memory region for accessing FicTrac's state
+    shmem = mmap.mmap(-1, ctypes.sizeof(fictrac.FicTracState), "FicTracStateSHMEM")
+
+    warnings.filterwarnings("ignore", category=matplotlib.cbook.mplDeprecation)
+    plt.ion()
+
+    fig = plt.figure()
+    fig.canvas.set_window_title('traces: fictrac')
+
+    # Number of fields to display
+    num_channels = len(fictrac_state_fields)
+
+    # Axes limits for each field
+    field_ax_limits = {'speed': (0, .03),
+                       'direction': (0, 2*np.pi),
+                       'heading': (0, 2*np.pi),
+                       'heading_diff': (0, 0.261799),
+                       'del_rot_error': (0, 15000),
+                       'del_rot_cam_vec': (-0.025, 0.025)}
+
+    # Setup a queue for caching the historical data received so we can plot history of samples up to
+    # some N
+    data_history = deque([FicTracState() for i in range(num_history)], maxlen=num_history)
+
+    plot_data = np.zeros((num_history, num_channels))
+
+    # We want to create a subplot for each channel
+    axes = []
+    backgrounds = []
+    point_sets = []
+    for chn in range(1, num_channels + 1):
+        field_name = fictrac_state_fields[chn-1]
+        ax = fig.add_subplot(num_channels, 1, chn)
+        ax.set_title(field_name)
+        backgrounds.append(fig.canvas.copy_from_bbox(ax.bbox))  # cache the background
+        ax.axis([0, num_history, field_ax_limits[field_name][0], field_ax_limits[field_name][1]])
+        axes.append(ax)
+        point_sets.append(ax.plot(np.arange(num_history), plot_data)[0])  # init plot content
+
+    plt.show()
+    plt.draw()
+    fig.canvas.start_event_loop(0.001)  # otherwise plot freezes after 3-4 iterations
+
+    RUN = True
+    data = fictrac.FicTracState.from_buffer(shmem)
+    first_frame_count = data.frame_cnt
+    old_frame_count = data.frame_cnt
+    while flyvr_state.is_running_well():
+        new_frame_count = data.frame_cnt
+
+        if old_frame_count != new_frame_count:
+
+            # Copy the current state.
+            data_copy = FicTracState()
+            ctypes.pointer(data_copy)[0] = data
+
+            # Append to the history
+            data_history.append(data_copy)
+
+            # Turned the queued chunks into a flat array
+            sample_i = 0
+            last_d = data_history[0]
+            for d in data_history:
+                chan_i = 0
+                for field in fictrac_state_fields:
+                    if field.endswith('_diff'):
+                        real_field = field.replace('_diff', '')
+
+                        if real_field in ['heading', 'direction']:
+                            plot_data[sample_i, chan_i] = abs(angle_diff(getattr(d, real_field), getattr(last_d, real_field)))
+                        else:
+                            plot_data[sample_i, chan_i] = getattr(d, real_field) - getattr(last_d, real_field)
+                    elif field.endswith('vec'):
+                        plot_data[sample_i, chan_i] = getattr(d, field)[1]
+                    else:
+                        plot_data[sample_i,chan_i] = getattr(d, field)
+                    chan_i = chan_i + 1
+                last_d = d
+                sample_i = sample_i + 1
+
+            for chn in range(num_channels):
+                fig.canvas.restore_region(backgrounds[chn])         # restore background
+                point_sets[chn].set_data(np.arange(num_history), plot_data[:,chn])
+                axes[chn].draw_artist(point_sets[chn])              # redraw just the points
+                #fig.canvas.blit(axes[chn].bbox)                    # fill in the axes rectangle
+
+            fig.canvas.draw()
+            fig.canvas.flush_events()
+
+    # clean up
+    plt.close(fig)