diff --git a/src/vidxp/search_fusion.py b/src/vidxp/search_fusion.py index b7e4184..8889798 100644 --- a/src/vidxp/search_fusion.py +++ b/src/vidxp/search_fusion.py @@ -33,7 +33,7 @@ def _query_id( return "fused:" + hashlib.sha256(identity.encode("utf-8")).hexdigest() -def _connected_components( +def _shared_overlap_components( hits: tuple[SearchHit, ...], ) -> list[list[SearchHit]]: ordered = sorted( @@ -50,17 +50,21 @@ def _connected_components( components: list[list[SearchHit]] = [] current: list[SearchHit] = [] current_media: str | None = None - current_end = 0.0 + current_overlap_end = 0.0 for hit in ordered: - if not current or hit.media_id != current_media or hit.start > current_end: + if ( + not current + or hit.media_id != current_media + or hit.start > current_overlap_end + ): if current: components.append(current) current = [hit] current_media = hit.media_id - current_end = hit.end + current_overlap_end = hit.end else: current.append(hit) - current_end = max(current_end, hit.end) + current_overlap_end = min(current_overlap_end, hit.end) if current: components.append(current) return components @@ -137,7 +141,7 @@ def fuse_search_results( ordered_results = tuple(by_modality[modality] for modality in searched_modalities) flattened = tuple(hit for result in ordered_results for hit in result.hits) candidates = [] - for hits in _connected_components(flattened): + for hits in _shared_overlap_components(flattened): ordered_hits = tuple( sorted( hits, @@ -152,8 +156,8 @@ def fuse_search_results( { "score": _score(hits), "media_id": hits[0].media_id, - "start": min(hit.start for hit in hits), - "end": max(hit.end for hit in hits), + "start": max(hit.start for hit in hits), + "end": min(hit.end for hit in hits), "modalities": tuple(sorted({hit.modality for hit in hits})), "hits": ordered_hits, } diff --git a/tests/test_search_fusion.py b/tests/test_search_fusion.py index 622e021..5611303 100644 --- a/tests/test_search_fusion.py +++ b/tests/test_search_fusion.py @@ -57,8 +57,8 @@ def test_rrf_counts_only_the_best_rank_per_modality_in_a_moment(self): moment = result.moments[0] self.assertAlmostEqual(moment.score, 2 / (RRF_RANK_CONSTANT + 1)) self.assertEqual(len(moment.hits), 3) - self.assertEqual(moment.start, 1) - self.assertEqual(moment.end, 4) + self.assertEqual(moment.start, 2.5) + self.assertEqual(moment.end, 3.0) def test_result_order_does_not_change_fusion_identity_or_output(self): scene = SearchResult( @@ -96,9 +96,7 @@ def test_rewritten_atomic_query_identity_changes_fused_identity(self): modality="scene", hits=(hit("scene", 1, 1, 2, "scene:1"),), ) - rewritten = original.model_copy( - update={"query_id": "scene:rewritten"} - ) + rewritten = original.model_copy(update={"query_id": "scene:rewritten"}) arguments = { "query": "Where is the taxi?", "requested_modalities": ("scene",), @@ -109,6 +107,69 @@ def test_rewritten_atomic_query_identity_changes_fused_identity(self): self.assertNotEqual(first.query_id, second.query_id) + def test_bridging_hit_does_not_merge_separate_moments(self): + hit_a = hit("scene", 1, 10.0, 12.0, "scene:a") + hit_b = hit("speech", 1, 11.0, 25.0, "speech:b") + hit_c = hit("scene", 2, 24.0, 26.0, "scene:c") + + scene = SearchResult( + query_id="scene:q", + query="car", + modality="scene", + hits=(hit_a, hit_c), + ) + speech = SearchResult( + query_id="speech:q", + query="car", + modality="speech", + hits=(hit_b,), + ) + + result = fuse_search_results( + query="car", + requested_modalities=("scene", "speech"), + results=(scene, speech), + ) + + self.assertEqual(len(result.moments), 2) + moment_1, moment_2 = result.moments + self.assertEqual(moment_1.start, 11.0) + self.assertEqual(moment_1.end, 12.0) + self.assertIn("scene:a", [h.source_id for h in moment_1.hits]) + self.assertEqual(moment_2.start, 24.0) + self.assertEqual(moment_2.end, 26.0) + self.assertIn("scene:c", [h.source_id for h in moment_2.hits]) + self.assertNotIn("scene:c", [h.source_id for h in moment_1.hits]) + + def test_nearby_duplicate_hits_combine_into_one_moment(self): + hit_1 = hit("scene", 1, 1.0, 3.0, "scene:1") + hit_2 = hit("scene", 2, 2.0, 3.5, "scene:2") + hit_3 = hit("speech", 1, 2.2, 2.8, "speech:1") + + scene = SearchResult( + query_id="scene:q", + query="dog", + modality="scene", + hits=(hit_1, hit_2), + ) + speech = SearchResult( + query_id="speech:q", + query="dog", + modality="speech", + hits=(hit_3,), + ) + + result = fuse_search_results( + query="dog", + requested_modalities=("scene", "speech"), + results=(scene, speech), + ) + + self.assertEqual(len(result.moments), 1) + self.assertEqual(len(result.moments[0].hits), 3) + self.assertEqual(result.moments[0].start, 2.2) + self.assertEqual(result.moments[0].end, 2.8) + if __name__ == "__main__": unittest.main()