updated SpeedTrigger/FreezeTrigger and added example SpeedExperiment

Author: JensBlack

experiments/custom/experiments.py

@@ -11,7 +11,7 @@
 from functools import partial
 from collections import Counter
 from experiments.custom.stimulus_process import ClassicProtocolProcess, SimpleProtocolProcess,Timer, ExampleProtocolProcess
-from experiments.custom.triggers import ScreenTrigger, RegionTrigger, OutsideTrigger, DirectionTrigger
+from experiments.custom.triggers import ScreenTrigger, RegionTrigger, OutsideTrigger, DirectionTrigger, SpeedTrigger
 from utils.plotter import plot_triggers_response
 from utils.analysis import angle_between_vectors
 from experiments.custom.stimulation import show_visual_stim_img,laser_switch
@@ -137,6 +137,80 @@ def get_trial(self):
 EXP_COMPLETION = 10
 
 
+
+class SpeedExperiment:
+    """
+    Simple class to contain all of the experiment properties
+    Uses multiprocess to ensure the best possible performance and
+        to showcase that it is possible to work with any type of equipment, even timer-dependent
+    """
+    def __init__(self):
+        self.experiment_finished = False
+        self._threshold = 10
+        self._event = None
+        self._current_trial = None
+        self._event_count = 0
+        self._trigger = SpeedTrigger(threshold = self._threshold,bodypart= 'tailroot', timewindow_len= 5)
+        self._exp_timer = Timer(600)
+
+    def check_skeleton(self, frame, skeleton):
+        """
+        Checking each passed animal skeleton for a pre-defined set of conditions
+        Outputting the visual representation, if exist
+        Advancing trials according to inherent logic of an experiment
+        :param frame: frame, on which animal skeleton was found
+        :param skeleton: skeleton, consisting of multiple joints of an animal
+        """
+        self.check_exp_timer()  # checking if experiment is still on
+
+        if not self.experiment_finished:
+            result, response = self._trigger.check_skeleton(skeleton=skeleton)
+            plot_triggers_response(frame, response)
+            if result:
+                laser_switch(True)
+                self._event_count += 1
+                print(self._event_count)
+                print('Light on')
+
+            else:
+                laser_switch(False)
+                print('Light off')
+
+            return result, response
+
+
+    def check_exp_timer(self):
+        """
+        Checking the experiment timer
+        """
+        if not self._exp_timer.check_timer():
+            print("Experiment is finished")
+            print("Time ran out.")
+            self.stop_experiment()
+
+    def start_experiment(self):
+        """
+        Start the experiment
+        """
+        if not self.experiment_finished:
+            self._exp_timer.start()
+
+    def stop_experiment(self):
+        """
+        Stop the experiment and reset the timer
+        """
+        self.experiment_finished = True
+        print('Experiment completed!')
+        self._exp_timer.reset()
+        # don't forget to stop the laser for safety!
+        laser_switch(False)
+
+    def get_trial(self):
+        """
+        Check which trial is going on right now
+        """
+        return self._current_trial
+
 class FirstExperiment:
     def __init__(self):
         self.experiment_finished = False

experiments/custom/triggers.py

@@ -9,7 +9,8 @@
 
 from utils.analysis import angle_between_vectors, calculate_distance, EllipseROI, RectangleROI
 from utils.configloader import RESOLUTION
-
+from collections import deque
+import numpy as np
 """Single posture triggers"""
 
 class HeaddirectionROITrigger:
@@ -333,15 +334,20 @@ def check_skeleton(self, skeleton: dict):
 
 class FreezeTrigger:
     """
-    Trigger to check if animal is in freezing state
+    Trigger to check if animal is moving below a certain speed
     """
-    def __init__(self, threshold: int, debug: bool = False):
+    def __init__(self, threshold: int, bodypart: str, timewindow_len:int = 2,  debug: bool = False):
         """
         Initializing trigger with given threshold
         :param threshold: int in pixel how much of a movement does not count
-        For example threshold of 5 would mean that all movements less then 5 pixels would be ignored
+        :param bodypart: str of body part in skeleton used for speed calculation
+        For example threshold of 5 would mean that all movements more then 5 pixels in the last timewindow length frames
+        would be ignored
         """
+        self._bodypart = bodypart
         self._threshold = threshold
+        self._timewindow_len = timewindow_len
+        self._timewindow = deque(maxlen= timewindow_len)
         self._skeleton = None
         self._debug = debug  # not used in this trigger
 
@@ -354,45 +360,52 @@ def check_skeleton(self, skeleton: dict):
         """
         # choosing a point to draw near the skeleton
         org_point = skeleton[list(skeleton.keys())[0]]
-        joint_moved = []
+        joint_moved = 0
+
         if self._skeleton is None:
             result = False
-            text = 'Not freezing'
+            text = '...'
             self._skeleton = skeleton
         else:
-            for joint in skeleton:
-                joint_travel = calculate_distance(skeleton[joint], self._skeleton[joint])
-                joint_moved.append(abs(joint_travel) <= self._threshold)
-            if all(joint_moved):
+            joint_travel = calculate_distance(skeleton[self._bodypart], self._skeleton[self._bodypart])
+            print(joint_travel)
+            self._timewindow.append(joint_travel)
+            print(self._timewindow)
+            if len(self._timewindow) == self._timewindow_len:
+                joint_moved = np.sum(self._timewindow)
+            print(joint_moved)
+
+            if abs(joint_moved) <= self._threshold:
                 result = True
                 text = 'Freezing'
             else:
                 result = False
-                text = 'Not freezing'
-            self._skeleton = skeleton
-
+                text = 'Not Freezing'
+        self._skeleton = skeleton
         color = (0, 255, 0) if result else (0, 0, 255)
         response_body = {'plot': {'text': dict(text=text,
                                                org=org_point,
                                                color=color)}}
         response = (result, response_body)
-        return response
 
+        return response
 
 class SpeedTrigger:
     """
     Trigger to check if animal is moving above a certain speed
     """
-    def __init__(self, threshold: int, bodypart: str = 'any', debug: bool = False):
+    def __init__(self, threshold: int, bodypart: str, timewindow_len:int = 2,  debug: bool = False):
         """
         Initializing trigger with given threshold
         :param threshold: int in pixel how much of a movement does not count
-        :param bodypart: str or list of str, bodypart or list of bodyparts in skeleton to use for trigger,
-         if "any" will check if any bodypart reaches treshold; default "any"
-        For example threshold of 5 would mean that all movements less then 5 pixels would be ignored
+        :param bodypart: str of body part in skeleton used for speed calculation
+        For example threshold of 5 would mean that all movements less then 5 pixels in the last timewindow length frames
+        would be ignored
         """
         self._bodypart = bodypart
         self._threshold = threshold
+        self._timewindow_len = timewindow_len
+        self._timewindow = deque(maxlen= timewindow_len)
         self._skeleton = None
         self._debug = debug  # not used in this trigger
 
@@ -405,36 +418,32 @@ def check_skeleton(self, skeleton: dict):
         """
         # choosing a point to draw near the skeleton
         org_point = skeleton[list(skeleton.keys())[0]]
-        joint_moved = []
+        joint_moved = 0
+
         if self._skeleton is None:
             result = False
-            text = 'First frame'
+            text = '...'
             self._skeleton = skeleton
         else:
-            if self._bodypart is "any":
-                for joint in skeleton:
-                    joint_travel = calculate_distance(skeleton[joint], self._skeleton[joint])
-                    joint_moved.append(abs(joint_travel) >= self._threshold)
-
-            elif isinstance(self._bodypart, list):
-                for joint in self._bodypart:
-                    joint_travel = calculate_distance(skeleton[joint], self._skeleton[joint])
-                    joint_moved.append(abs(joint_travel) >= self._threshold)
-            else:
-                joint_travel = calculate_distance(skeleton[self._bodypart], self._skeleton[self._bodypart])
-                joint_moved.append(abs(joint_travel) >= self._threshold)
-
-            if all(joint_moved):
+            joint_travel = calculate_distance(skeleton[self._bodypart], self._skeleton[self._bodypart])
+            print(joint_travel)
+            self._timewindow.append(joint_travel)
+            print(self._timewindow)
+            if len(self._timewindow) == self._timewindow_len:
+                joint_moved = np.sum(self._timewindow)
+            print(joint_moved)
+
+            if abs(joint_moved) >= self._threshold:
                 result = True
                 text = 'Running'
             else:
                 result = False
-                text = ''
-            self._skeleton = skeleton
-
+                text = 'Not Running'
+        self._skeleton = skeleton
         color = (0, 255, 0) if result else (0, 0, 255)
         response_body = {'plot': {'text': dict(text=text,
                                                org=org_point,
                                                color=color)}}
         response = (result, response_body)
+
         return response