Merge pull request #1003 from int-brain-lab/lpqc

PostLP task

Author: mayofaulkner

brainbox/behavior/dlc.py

@@ -44,16 +44,20 @@ def insert_idx(array, values):
     return idx
 
 
+def valid_feature(x: str):
+    if x.endswith('_x') or x.endswith('_y') or x.endswith('_likelihood'):
+        return True
+    return False
+
+
 def likelihood_threshold(dlc, threshold=0.9):
-    """
-    Set dlc points with likelihood less than threshold to nan.
+    """Set dlc points with likelihood less than threshold to nan.
 
-    FIXME Add unit test.
     :param dlc: dlc pqt object
     :param threshold: likelihood threshold
     :return:
     """
-    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc.keys()])
+    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc.keys() if valid_feature(x)])
     for feat in features:
         nan_fill = dlc[f'{feat}_likelihood'] < threshold
         dlc.loc[nan_fill, (f'{feat}_x', f'{feat}_y')] = np.nan
@@ -268,7 +272,7 @@ def plot_trace_on_frame(frame, dlc_df, cam):
     # Threshold the dlc traces
     dlc_df = likelihood_threshold(dlc_df)
     # Features without tube
-    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc_df.keys() if 'tube' not in x])
+    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc_df.keys() if valid_feature(x) and 'tube' not in x])
     # Normalize the number of points across cameras
     dlc_df_norm = pd.DataFrame()
     for feat in features:

brainbox/io/one.py

@@ -1551,7 +1551,7 @@ def load_wheel(self, fs=1000, corner_frequency=20, order=8, collection=None):
         self.wheel = self.wheel.apply(np.float32)
         self.data_info.loc[self.data_info['name'] == 'wheel', 'is_loaded'] = True
 
-    def load_pose(self, likelihood_thr=0.9, views=['left', 'right', 'body']):
+    def load_pose(self, likelihood_thr=0.9, views=['left', 'right', 'body'], tracker='dlc'):
         """
         Function to load the pose estimation results (DLC) into SessionLoader.pose. SessionLoader.pose is a
         dictionary where keys are the names of the cameras for which pose data is loaded, and values are pandas
@@ -1565,13 +1565,17 @@ def load_pose(self, likelihood_thr=0.9, views=['left', 'right', 'body']):
             likelihood_thr=1. Default is 0.9
         views: list
             List of camera views for which to try and load data. Possible options are {'left', 'right', 'body'}
+        tracker : str
+            Tracking algorithm to load pose estimates from. Possible options are {'dlc', 'lightningPose'}
         """
         # empty the dictionary so that if one loads only one view, after having loaded several, the others don't linger
+        tracker = 'lightningPose' if tracker in ['lp', 'litpose'] else tracker
         self.pose = {}
         for view in views:
-            pose_raw = self.one.load_object(self.eid, f'{view}Camera', attribute=['dlc', 'times'], revision=self.revision or None)
+            pose_raw = self.one.load_object(
+                self.eid, f'{view}Camera', attribute=[tracker, 'times'], revision=self.revision or None)
             # Double check if video timestamps are correct length or can be fixed
-            times_fixed, dlc = self._check_video_timestamps(view, pose_raw['times'], pose_raw['dlc'])
+            times_fixed, dlc = self._check_video_timestamps(view, pose_raw['times'], pose_raw[tracker])
             self.pose[f'{view}Camera'] = likelihood_threshold(dlc, likelihood_thr)
             self.pose[f'{view}Camera'].insert(0, 'times', times_fixed)
             self.data_info.loc[self.data_info['name'] == 'pose', 'is_loaded'] = True

brainbox/tests/test_behavior.py

@@ -3,17 +3,153 @@
 from unittest import mock
 from functools import partial
 import numpy as np
+import pandas as pd
 import pickle
 import copy
 
 from iblutil.util import Bunch
 from one.api import ONE
 
+import brainbox.behavior.dlc as dlc
 import brainbox.behavior.wheel as wheel
 import brainbox.behavior.training as train
 from ibllib.tests import TEST_DB
 
 
+class TestDLC(unittest.TestCase):
+
+    def setUp(self):
+        pass
+
+    def test_plt_window(self):
+        """Test for brainbox.behavior.dlc.plt_window"""
+        window_lag = -0.5
+        window_len = 2
+        x = 10
+        beg, end = dlc.plt_window(x)
+        self.assertTrue(beg == x + window_lag, msg='Unexpected beg time')
+        self.assertTrue(end == x + window_len, msg='Unexpected beg time')
+
+    def test_insert_idx(self):
+        """Test for brainbox.behavior.dlc.insert_idx"""
+
+        # Test basic functionality with simple arrays
+        array = np.array([1, 3, 5, 7, 9])
+        values = np.array([2, 6])
+        result = dlc.insert_idx(array, values)
+        expected = np.array([1, 3])
+        np.testing.assert_array_equal(result, expected)
+
+        # Test when values exactly match array elements
+        array = np.array([1, 3, 5, 7, 9])
+        values = np.array([3, 7])
+        result = dlc.insert_idx(array, values)
+        expected = np.array([1, 3])
+        np.testing.assert_array_equal(result, expected)
+
+        # Test values at array boundaries
+        array = np.array([2, 4, 6, 8, 10])
+        values = np.array([1, 11])  # Below first, above last
+        result = dlc.insert_idx(array, values)
+        expected = np.array([0, 4])
+        np.testing.assert_array_equal(result, expected)
+
+        # Test with single value to insert
+        array = np.array([1, 3, 5, 7, 9])
+        values = np.array([4])
+        result = dlc.insert_idx(array, values)
+        expected = np.array([2])
+        np.testing.assert_array_equal(result, expected)
+
+        # Test that ValueError is raised when all values map to index 0
+        array = np.array([10, 20, 30])
+        values = np.array([1, 2, 3])  # All much closer to first element
+        with self.assertRaises(ValueError) as context:
+            dlc.insert_idx(array, values)
+        self.assertIn('Something is wrong, all values to insert are outside of the array', str(context.exception))
+
+        # Test with negative values in array and values
+        array = np.array([-5, -2, 1, 4])
+        values = np.array([-3, 0])
+        result = dlc.insert_idx(array, values)
+        expected = np.array([1, 2])
+        np.testing.assert_array_equal(result, expected)
+
+    def test_valid_feature(self):
+        """Test for brainbox.behavior.dlc.valid_feature"""
+
+        valid = dlc.valid_feature('test_x')
+        self.assertTrue(valid, msg='strings ending in "x" should be valid')
+
+        valid = dlc.valid_feature('test_y')
+        self.assertTrue(valid, msg='strings ending in "y" should be valid')
+
+        valid = dlc.valid_feature('test_likelihood')
+        self.assertTrue(valid, msg='strings ending in "likelihood" should be valid')
+
+        valid = dlc.valid_feature('test_l')
+        self.assertTrue(not valid, msg='only strings ending in "_x", "_y" or "_likelihood" are valid')
+
+        valid = dlc.valid_feature('testlikelihood')
+        self.assertTrue(not valid, msg='only strings ending in "_x", "_y" or "_likelihood" are valid')
+
+    def test_likelihood_threshold(self):
+        """Test for brainbox.behavior.dlc.likelihood_threshold"""
+
+        dlc_data = pd.DataFrame({
+            'nose_x': [10.0, 20.0, 30.0, 40.0, 50.0],
+            'nose_y': [15.0, 25.0, 35.0, 45.0, 55.0],
+            'nose_likelihood': [0.95, 0.85, 0.92, 0.88, 0.91],
+            'ear_x': [12.0, 22.0, 32.0, 42.0, 52.0],
+            'ear_y': [17.0, 27.0, 37.0, 47.0, 57.0],
+            'ear_likelihood': [0.99, 0.75, 0.80, 0.95, 0.60],
+            'tail_x': [5.0, 15.0, 25.0, 35.0, 45.0],
+            'tail_y': [8.0, 18.0, 28.0, 38.0, 48.0],
+            'tail_likelihood': [0.70, 0.95, 0.85, 0.92, 0.88]
+        })
+
+        # Test with default threshold of 0.9
+        result = dlc.likelihood_threshold(dlc_data.copy())
+        nans = {
+            'nose': np.array([1, 3]),
+            'ear': np.array([1, 2, 4]),
+            'tail': np.array([0, 2, 4]),
+        }
+        for kp in nans.keys():
+            for idx in range(5):
+                if len(np.where(nans[kp] == idx)[0]) > 0:
+                    self.assertTrue(pd.isna(result.loc[idx, f'{kp}_x']))
+                    self.assertTrue(pd.isna(result.loc[idx, f'{kp}_y']))
+                else:
+                    self.assertFalse(pd.isna(result.loc[idx, f'{kp}_x']))
+                    self.assertFalse(pd.isna(result.loc[idx, f'{kp}_y']))
+
+        # Test with custom threshold
+        result = dlc.likelihood_threshold(dlc_data.copy(), threshold=0.8)
+        nans = {
+            'nose': np.array([]),
+            'ear': np.array([1, 4]),
+            'tail': np.array([0]),
+        }
+        for kp in nans.keys():
+            for idx in range(5):
+                if len(np.where(nans[kp] == idx)[0]) > 0:
+                    self.assertTrue(pd.isna(result.loc[idx, f'{kp}_x']))
+                    self.assertTrue(pd.isna(result.loc[idx, f'{kp}_y']))
+                else:
+                    self.assertFalse(pd.isna(result.loc[idx, f'{kp}_x']))
+                    self.assertFalse(pd.isna(result.loc[idx, f'{kp}_y']))
+
+        # Test with dataframe containing no valid features
+        invalid_data = pd.DataFrame({
+            'random_col1': [1, 2, 3],
+            'random_col2': [4, 5, 6]
+        })
+        result = dlc.likelihood_threshold(invalid_data)
+        # Should return unchanged dataframe
+        pd.testing.assert_frame_equal(result, invalid_data)
+
+
 class TestWheel(unittest.TestCase):
 
     def setUp(self):

ibllib/oneibl/patcher.py

@@ -682,7 +682,7 @@ def patch_dataset(self, file_list, dry=False, ftp=False, force=False, **kwargs):
             _logger.error(f'Files: {", ".join([f.name for f in file_list])} already exist, to overwrite set force=True')
             return
 
-        response = super().patch_dataset(file_list, dry=dry, repository=self.s3_repo, ftp=False, **kwargs)
+        response = super().patch_dataset(file_list, dry=dry, repository=self.s3_repo, ftp=False, force=force, **kwargs)
         # TODO in an ideal case the flatiron filerecord won't be altered when we register this dataset. This requires
         # changing the the alyx.data.register_view
         for ds in response: