Added digging to the list of detected behaviours by deepof.data.coordinates.rule_based_annotation()

deepof/pose_utils.py

@@ -180,20 +180,19 @@ def huddle(
     tol_speed: float,
     animal_id: str = "",
 ) -> np.array:
-    """Returns true when the mouse is huddling using simple rules. (!!!) Designed to
-    work with deepof's default DLC mice models; not guaranteed to work otherwise.
+    """Returns true when the mouse is huddling using simple rules.
 
-        Parameters:
-            - pos_dframe (pandas.DataFrame): position of body parts over time
-            - speed_dframe (pandas.DataFrame): speed of body parts over time
-            - tol_forward (float): Maximum tolerated distance between ears and
-            forward limbs
-            - tol_rear (float): Maximum tolerated average distance between spine
-            body parts
-            - tol_speed (float): Maximum tolerated speed for the center of the mouse
+    Parameters:
+        - pos_dframe (pandas.DataFrame): position of body parts over time
+        - speed_dframe (pandas.DataFrame): speed of body parts over time
+        - tol_forward (float): Maximum tolerated distance between ears and
+        forward limbs
+        - tol_rear (float): Maximum tolerated average distance between spine
+        body parts
+        - tol_speed (float): Maximum tolerated speed for the center of the mouse
 
-        Returns:
-            hudd (np.array): True if the animal is huddling, False otherwise
+    Returns:
+        hudd (np.array): True if the animal is huddling, False otherwise
     """
 
     if animal_id != "":
@@ -220,15 +219,37 @@ def huddle(
 
 
 def dig(
-    pos_dframe: pd.DataFrame,
     speed_dframe: pd.DataFrame,
     likelihood_dframe: pd.DataFrame,
-    tol_nose_speed: float,
     tol_speed: float,
     tol_likelihood: float,
     animal_id: str = "",
 ):
-    pass
+    """Returns true when the mouse is digging using simple rules.
+
+    Parameters:
+        - speed_dframe (pandas.DataFrame): speed of body parts over time
+        - likelihood_dframe (pandas.DataFrame): likelihood of body part tracker over time,
+        as directly obtained from DeepLabCut
+        - tol_nose_speed (float): Maximum tolerated average distance between spine
+        body parts
+        - tol_speed (float): Maximum tolerated speed for the center of the mouse
+        - tol_likelihood (float): Maximum tolerated likelihood for the nose (if the animal
+        is digging, the nose is momentarily occluded).
+
+    Returns:
+        dig (np.array): True if the animal is digging, False otherwise
+    """
+
+    if animal_id != "":
+        animal_id += "_"
+
+    speed = speed_dframe[animal_id + "Center"] < tol_speed
+    nose_speed = speed_dframe[animal_id + "Center"] < speed_dframe[animal_id + "Nose"]
+    likelihood = likelihood_dframe[animal_id + "Nose"] < tol_likelihood
+    dig = speed & nose_speed & likelihood
+
+    return dig
 
 
 def sniff(