Coverage for cafl4ds/filters/dedup.py: 100%

1"""SemDeDup-style near-duplicate removal — an admission knob (Abbas et al. 2023). 

2 

3SemDeDup deduplicates a dataset by embedding it, clustering, and within each cluster dropping 

4points whose cosine similarity to a kept neighbour exceeds a threshold — keeping one 

5representative per near-duplicate group. Here we specialise the *within-cluster* step to a 

6single streaming batch: the batch is small, so no k-means is needed — we encode it with the 

7**current** model, then greedily keep points and drop any later point too similar (cosine) to 

8an already-kept one. This trims the redundancy a correlated stream pours in (many near-identical 

9consecutive frames), so the SSL update sees a more diverse diet. 

10 

11This is a *label-free* selection heuristic over the live embedding (the fast edge); it never 

12touches labels or the SSL loss. Composes with the other admission knobs and, upstream of a 

13:class:`~cafl4ds.filters.reservoir.ReservoirReplay`, gives baseline **B1.5** (reservoir + dedup). 

14""" 

15 

16from __future__ import annotations 

17 

18import torch 

19 

20from cafl4ds.data.streams import StreamBatch 

21from cafl4ds.filters.base import Filter, FilterContext 

22 

23 

24def semantic_dedup_keep(embeddings: torch.Tensor, threshold: float) -> torch.Tensor: 

25 """Greedy near-duplicate removal in embedding space (the SemDeDup within-cluster step). 

26 

27 Iterates the rows in order, keeping a point unless its cosine similarity to some 

28 already-kept point exceeds ``threshold`` (in which case it is a near-duplicate and 

29 dropped). Keeping the *earliest* of each near-duplicate group preserves stream order. 

30 

31 Args: 

32 embeddings: Row-wise embeddings ``[B, d]`` (need not be normalized). 

33 threshold: Cosine-similarity above which two points are treated as near-duplicates, 

34 in ``[-1, 1]``. Higher = more permissive (only near-identical points are dropped); 

35 ``>= 1.0`` keeps everything. 

36 

37 Returns: 

38 A 1-D ``long`` tensor of kept row indices, in ascending (stream) order. 

39 """ 

40 n = embeddings.shape[0] 

41 if n <= 1: 

42 return torch.arange(n, dtype=torch.long, device=embeddings.device) 

43 normed = torch.nn.functional.normalize(embeddings, dim=1) 

44 kept: list[int] = [] 

45 for i in range(n): 

46 if not kept: 

47 kept.append(i) 

48 continue 

49 sims = normed[i] @ normed[torch.tensor(kept, device=embeddings.device)].T # [len(kept)] 

50 if float(sims.max()) <= threshold: 

51 kept.append(i) 

52 return torch.tensor(kept, dtype=torch.long, device=embeddings.device) 

53 

54 

55class SemDeDup(Filter): 

56 """Admission knob: drop near-duplicate frames by cosine similarity in embedding space.""" 

57 

58 def __init__(self, threshold: float = 0.9) -> None: 

59 """Configure the deduplicator. 

60 

61 Args: 

62 threshold: Cosine-similarity above which a frame is dropped as a near-duplicate of 

63 an already-kept frame. ``0.9`` keeps clearly-distinct frames and removes only 

64 the near-identical ones; raise toward ``1.0`` to dedup less aggressively. 

65 """ 

66 self.threshold = threshold 

67 

68 def select(self, batch: StreamBatch, ctx: FilterContext) -> torch.Tensor: 

69 """Encode the batch with the live model and return its near-duplicate-free subset. 

70 

71 Args: 

72 batch: The incoming label-free stream batch. 

73 ctx: The live filter context; its ``method`` supplies the embedding (fast edge). 

74 

75 Returns: 

76 The kept images ``[K, C, H, W]`` (``K <= B``), earliest of each duplicate group. 

77 """ 

78 embeddings = ctx.method.encode(batch.images) # [B, d], no grad 

79 keep = semantic_dedup_keep(embeddings, self.threshold) 

80 return batch.images[keep]