Merge pull request #107 from KumarLabJax/bugfix/empty-detection-edge-cases

Fix edge cases for empty detection sets in vectorized matching

Author: bergsalex

src/mouse_tracking/matching/core.py

@@ -601,9 +601,14 @@ def __init__(self, observations: list[list[Detection]]):
 
         # Metadata
         self._num_frames = len(observations)
-        self._median_observation = int(np.median([len(x) for x in observations]))
-        # Add 0.5 to do proper rounding with int cast
-        self._avg_observation = int(np.mean([len(x) for x in observations]) + 0.5)
+        # Handle empty observation list edge case
+        if len(observations) == 0:
+            self._median_observation = 0
+            self._avg_observation = 0
+        else:
+            self._median_observation = int(np.median([len(x) for x in observations]))
+            # Add 0.5 to do proper rounding with int cast
+            self._avg_observation = int(np.mean([len(x) for x in observations]) + 0.5)
         self._tracklet_gen_method = None
         self._tracklet_stitch_method = None
 
@@ -786,6 +791,10 @@ def get_embed_centers(self):
         longterm_ids = np.asarray(list(set(self._stitch_translation.values())))
         longterm_ids = longterm_ids[longterm_ids != 0]
 
+        # Handle edge case where all longterm IDs are 0 (filtered out)
+        if len(longterm_ids) == 0:
+            return np.zeros([0, embedding_shape[0]])
+
         # To calculate an average for merged tracklets, we weight by number of frames
         longterm_data = {}
         for cur_tracklet in self._tracklets:
@@ -841,8 +850,12 @@ def _make_tracklets(self, include_unassigned: bool = True):
         for tracklet_id, observation_list in tracklet_dict.items():
             tracklet_list.append(Tracklet(tracklet_id, observation_list))
 
-        if include_unassigned:
-            cur_tracklet_id = np.max(np.asarray(list(tracklet_dict.keys())))
+        if include_unassigned and len(unmatched_observations) > 0:
+            # Handle edge case where tracklet_dict is empty
+            if len(tracklet_dict) > 0:
+                cur_tracklet_id = np.max(np.asarray(list(tracklet_dict.keys())))
+            else:
+                cur_tracklet_id = 0
             for cur_observation in unmatched_observations:
                 tracklet_list.append(Tracklet(int(cur_tracklet_id), [cur_observation]))
                 cur_tracklet_id += 1
@@ -1079,35 +1092,42 @@ def generate_greedy_tracklets(
         if num_threads > 1:
             self._start_pool(num_threads)
 
-        # Main loop to cycle over greedy matching.
-        # Each match problem is posed as a bipartite graph between sequential frames
-        for frame in np.arange(len(self._observations) - 1) + 1:
-            # Cache the segmentation and rotation data
-            for obs in self._observations[frame - 1]:
-                obs.cache()
-            for obs in self._observations[frame]:
-                obs.cache()
-            # Calculate cost and greedily match
-            match_costs = self._calculate_costs(frame - 1, frame, rotate_pose)
-            match_costs = np.ma.array(match_costs, fill_value=max_cost, mask=False)
-            matches = {}
-            while np.any(~match_costs.mask) and np.any(match_costs.filled() < max_cost):
-                next_best = np.unravel_index(np.argmin(match_costs), match_costs.shape)
-                matches[next_best[1]] = prev_matches[next_best[0]]
-                match_costs.mask[next_best[0], :] = True
-                match_costs.mask[:, next_best[1]] = True
-            # Fill any unmatched observations
-            for j in range(len(self._observations[frame])):
-                if j not in matches:
-                    matches[j] = cur_tracklet_id
-                    cur_tracklet_id += 1
-            frame_dict[frame] = matches
-            # Cleanup for next loop iteration
-            for cur_obs in self._observations[frame - 1]:
-                cur_obs.clear_cache()
-            prev_matches = matches
-        if self._pool is not None:
-            self._kill_pool()
+        try:
+            # Main loop to cycle over greedy matching.
+            # Each match problem is posed as a bipartite graph between sequential frames
+            for frame in np.arange(len(self._observations) - 1) + 1:
+                # Cache the segmentation and rotation data
+                for obs in self._observations[frame - 1]:
+                    obs.cache()
+                for obs in self._observations[frame]:
+                    obs.cache()
+                # Calculate cost and greedily match
+                match_costs = self._calculate_costs(frame - 1, frame, rotate_pose)
+                match_costs = np.ma.array(match_costs, fill_value=max_cost, mask=False)
+                matches = {}
+                while np.any(~match_costs.mask) and np.any(
+                    match_costs.filled() < max_cost
+                ):
+                    next_best = np.unravel_index(
+                        np.argmin(match_costs), match_costs.shape
+                    )
+                    matches[next_best[1]] = prev_matches[next_best[0]]
+                    match_costs.mask[next_best[0], :] = True
+                    match_costs.mask[:, next_best[1]] = True
+                # Fill any unmatched observations
+                for j in range(len(self._observations[frame])):
+                    if j not in matches:
+                        matches[j] = cur_tracklet_id
+                        cur_tracklet_id += 1
+                frame_dict[frame] = matches
+                # Cleanup for next loop iteration
+                for cur_obs in self._observations[frame - 1]:
+                    cur_obs.clear_cache()
+                prev_matches = matches
+        finally:
+            # Ensure pool is always cleaned up, even if an exception occurs
+            if self._pool is not None:
+                self._kill_pool()
         # Final modification of internal state
         self._observation_id_dict = frame_dict
         self._tracklet_gen_method = "greedy"

src/mouse_tracking/matching/detection.py

@@ -98,8 +98,12 @@ def rotate_pose(
                 center: optional center of rotation. If not provided, the mean of non-tail keypoints are used as the center.
 
         Returns:
-                rotated keypoints
+                rotated keypoints, or None if points is None
         """
+        # Handle None input gracefully
+        if points is None:
+            return None
+
         points_valid = ~np.all(points == 0, axis=-1)
         # No points to rotate, just return original points.
         if np.all(~points_valid):
@@ -137,6 +141,9 @@ def embed_distance(embed_1, embed_2) -> float:
         Returns:
                 cosine distance between the embeddings
         """
+        # Check for None embeddings
+        if embed_1 is None or embed_2 is None:
+            return np.nan
         # Check for default embeddings
         if np.all(embed_1 == 0) or np.all(embed_2 == 0):
             return np.nan

src/mouse_tracking/matching/vectorized_features.py

@@ -135,6 +135,10 @@ def compute_vectorized_pose_distances(
     Returns:
             Distance matrix of shape (n1, n2) with mean pose distances
     """
+    # Handle edge case where either set has no detections
+    if features1.n_detections == 0 or features2.n_detections == 0:
+        return np.full((features1.n_detections, features2.n_detections), np.nan)
+
     poses1 = features1.poses  # Shape: (n1, 12, 2)
     poses2 = features2.poses  # Shape: (n2, 12, 2)
     valid1 = features1.valid_pose_masks  # Shape: (n1, 12)

tests/matching/core/video_observations/test_generate_greedy_tracklets.py

@@ -102,12 +102,14 @@ def test_generate_greedy_tracklets_no_observations(self):
         observations = [[] for _ in range(3)]  # All empty frames
         video_obs = VideoObservations(observations)
 
-        # TODO: This reveals a bug - _make_tracklets fails with empty tracklet_dict
-        # The _make_tracklets method tries to call np.max on empty array
-        with pytest.raises(
-            ValueError, match="zero-size array to reduction operation maximum"
-        ):
-            video_obs.generate_greedy_tracklets()
+        # Should handle empty frames gracefully
+        video_obs.generate_greedy_tracklets()
+
+        # Should have empty observation_id_dict and empty tracklets
+        assert video_obs._observation_id_dict is not None
+        assert video_obs._tracklet_gen_method == "greedy"
+        assert video_obs._tracklets is not None
+        assert len(video_obs._tracklets) == 0  # No tracklets for no observations
 
     def test_generate_greedy_tracklets_single_observation_per_frame(
         self, basic_detection
@@ -362,24 +364,33 @@ def test_generate_greedy_tracklets_pool_cleanup_on_exception(self, basic_detecti
 
         video_obs = VideoObservations(observations)
 
+        # Mock the pool object
+        mock_pool = MagicMock()
+
+        def mock_start_pool_impl(num_threads):
+            video_obs._pool = mock_pool
+
+        def mock_kill_pool_impl():
+            video_obs._pool = None
+
         with (
             patch.object(video_obs, "_start_pool") as mock_start_pool,
             patch.object(video_obs, "_kill_pool") as mock_kill_pool,
             patch.object(
                 video_obs, "_calculate_costs", side_effect=RuntimeError("Test error")
             ),
         ):
+            # Set up side effects so the mocks actually update _pool
+            mock_start_pool.side_effect = mock_start_pool_impl
+            mock_kill_pool.side_effect = mock_kill_pool_impl
+
             with pytest.raises(RuntimeError):
                 video_obs.generate_greedy_tracklets(num_threads=2)
 
             # Pool should be started
             mock_start_pool.assert_called_once()
-            # TODO: This reveals a bug - pool is not cleaned up on exception
-            # The generate_greedy_tracklets method doesn't use try/finally for cleanup
-            # Currently the pool is NOT cleaned up on exception
-            assert (
-                mock_kill_pool.call_count == 0
-            )  # Documents the current buggy behavior
+            # Pool should be cleaned up even though an exception occurred
+            mock_kill_pool.assert_called_once()
 
     def test_generate_greedy_tracklets_variable_observations_per_frame(
         self, basic_detection
@@ -443,10 +454,12 @@ def test_generate_greedy_tracklets_with_none_values(self, basic_detection):
 
         video_obs = VideoObservations(observations)
 
-        # TODO: This reveals a bug - rotate_pose doesn't handle None poses correctly
-        # The rotate_pose method assumes points is not None
-        with pytest.raises(TypeError, match="unsupported operand type"):
-            video_obs.generate_greedy_tracklets()
+        # Should handle None poses gracefully (using default costs)
+        video_obs.generate_greedy_tracklets(rotate_pose=True)
+
+        # Should complete without crashing
+        assert video_obs._tracklets is not None
+        assert video_obs._tracklet_gen_method == "greedy"
 
     def test_generate_greedy_tracklets_large_cost_matrix(self, basic_detection):
         """Test with larger cost matrices to ensure scalability."""
@@ -533,11 +546,15 @@ def test_generate_greedy_tracklets_deterministic_behavior(self, basic_detection)
 
     def test_generate_greedy_tracklets_empty_observation_list(self):
         """Test with empty observation list."""
-        # TODO: This reveals a bug - VideoObservations constructor can't handle empty lists
-        # The constructor tries to calculate median of empty list
-        with pytest.raises(ValueError, match="cannot convert float NaN to integer"):
-            observations = []
-            VideoObservations(observations)
+        # Should handle empty observation list gracefully
+        observations = []
+        video_obs = VideoObservations(observations)
+
+        # Verify attributes are set correctly
+        assert video_obs._num_frames == 0
+        assert video_obs._median_observation == 0
+        assert video_obs._avg_observation == 0
+        assert video_obs._observations == []
 
     def test_generate_greedy_tracklets_numerical_stability(self, basic_detection):
         """Test with edge cases that might cause numerical issues."""