From c3d4af94dc967fe505e85ba74c4d49bb2b963b5c Mon Sep 17 00:00:00 2001 From: Alexander Berger Date: Fri, 24 Oct 2025 11:20:07 -0400 Subject: [PATCH 1/4] Fix edge cases for empty detection sets in vectorized matching --- src/mouse_tracking/matching/core.py | 12 ++++++++++-- src/mouse_tracking/matching/vectorized_features.py | 4 ++++ .../test_generate_greedy_tracklets.py | 14 ++++++++------ 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/src/mouse_tracking/matching/core.py b/src/mouse_tracking/matching/core.py index 01d3c5f9..048d848b 100644 --- a/src/mouse_tracking/matching/core.py +++ b/src/mouse_tracking/matching/core.py @@ -786,6 +786,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 +845,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 diff --git a/src/mouse_tracking/matching/vectorized_features.py b/src/mouse_tracking/matching/vectorized_features.py index 526a2e11..2809d9aa 100644 --- a/src/mouse_tracking/matching/vectorized_features.py +++ b/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) diff --git a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py index 9a0bab62..294521fd 100644 --- a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py +++ b/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 From bdb3a788a975ba2f9cd0d05ddcb73694a0b77ed5 Mon Sep 17 00:00:00 2001 From: Alexander Berger Date: Fri, 24 Oct 2025 15:29:50 -0400 Subject: [PATCH 2/4] Additional defensive none pose handling --- src/mouse_tracking/matching/detection.py | 9 ++++++++- .../test_generate_greedy_tracklets.py | 10 ++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/mouse_tracking/matching/detection.py b/src/mouse_tracking/matching/detection.py index efd1a36a..ab1679ca 100644 --- a/src/mouse_tracking/matching/detection.py +++ b/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 diff --git a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py index 294521fd..b02db0f3 100644 --- a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py +++ b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py @@ -445,10 +445,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.""" From 18ee5d73575c24be5909016ddd07c8e533cdb584 Mon Sep 17 00:00:00 2001 From: Alexander Berger Date: Fri, 24 Oct 2025 15:38:39 -0400 Subject: [PATCH 3/4] Improve pool cleanup implementation --- src/mouse_tracking/matching/core.py | 65 ++++++++++--------- .../test_generate_greedy_tracklets.py | 21 ++++-- 2 files changed, 51 insertions(+), 35 deletions(-) diff --git a/src/mouse_tracking/matching/core.py b/src/mouse_tracking/matching/core.py index 048d848b..4b6f5fe2 100644 --- a/src/mouse_tracking/matching/core.py +++ b/src/mouse_tracking/matching/core.py @@ -1087,35 +1087,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" diff --git a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py index b02db0f3..9aee7040 100644 --- a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py +++ b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py @@ -364,6 +364,15 @@ 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, @@ -371,17 +380,17 @@ def test_generate_greedy_tracklets_pool_cleanup_on_exception(self, basic_detecti 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 From c9a81137a6b2149e44844613daf83b55b1f0e0a0 Mon Sep 17 00:00:00 2001 From: Alexander Berger Date: Fri, 24 Oct 2025 15:47:41 -0400 Subject: [PATCH 4/4] Improve VideoObservations constructor to avoid NaN conversion errors --- src/mouse_tracking/matching/core.py | 11 ++++++++--- .../test_generate_greedy_tracklets.py | 14 +++++++++----- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/mouse_tracking/matching/core.py b/src/mouse_tracking/matching/core.py index 4b6f5fe2..2d34d2d2 100644 --- a/src/mouse_tracking/matching/core.py +++ b/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 diff --git a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py index 9aee7040..93439b4c 100644 --- a/tests/matching/core/video_observations/test_generate_greedy_tracklets.py +++ b/tests/matching/core/video_observations/test_generate_greedy_tracklets.py @@ -546,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."""