Regroup the library without stopping the window
Pressing Regroup on a real library did not come back. Clustering 1,813 faces is the textbook agglomeration — compute every pairwise cosine, then repeatedly scan all live group pairs, score each with average link, and merge the best — and the scan is inside the loop. Each merge rescans every surviving pair, and each score is recomputed from scratch over every cross pair. Some 1.6 million pair scores per merge, some 700 merges to do. Three changes, none of which alter the answer. Only above-threshold pairs can ever matter. An average that reaches the threshold must have at least one term at or above it, so two groups with no qualifying pair between them can never merge — not now, and not after any sequence of merges, since merging only adds terms. The new `neighbours` module produces exactly that sparse list: 7,875 pairs rather than 1.6 million on the reference library. It also means the n^2 matrix is never materialised, so memory goes from O(n^2) to O(edges) — 2.5 GB to a few hundred KB at 25,000 faces. Merges cannot cross components, so the connected components of that graph are independent problems: four hundred small agglomerations instead of one large one. Average link is additive — sum(A u B, C) = sum(A, C) + sum(B, C) — so a merged group's scores follow by addition. Kept as running (sum, count) per adjacent pair, a score costs one division instead of a nested loop, and a heap with lazy invalidation replaces the rescan. Measured on the reference library: 0.28s, release, for all 1,813 faces. An exact ANN index was tried and removed, and neighbours.rs records why so it is not rediscovered as a good idea. IVF with a triangle-inequality bound is exact and prunes beautifully on synthetic clusters; on real embeddings it prunes *nothing* — 946 of 946 cell pairs survive. Median pair angle is 88.5 degrees and the merge threshold is 66.2, so the bound needs cells of radius under ~10 degrees, but two photographs of the same person sit 36-60 degrees apart. No ball-based partition of a 512-d near-orthogonal space can be tight enough. So the scan stayed exhaustive and got an unrolled dot product and its blocks spread across cores instead. Correctness is held by keeping the old implementation as an oracle: three tests run both engines over the same population — plain, under co-occurrence and anchor constraints, and with a size-weighted calibration — and assert the clusters are identical. Determinism is asserted at a size where the threaded path is in play. 62 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
+628
-110
@@ -24,11 +24,50 @@
|
||||
//! Single-link chains: one bad edge welds two identities together, and it is
|
||||
//! the documented way face clustering fails on families. Average link asks
|
||||
//! whether the *groups* are similar, which one outlier cannot force.
|
||||
//!
|
||||
//! # How it runs, and why the obvious way does not
|
||||
//!
|
||||
//! The first implementation of this was the textbook one: compute every
|
||||
//! pairwise cosine, then repeatedly scan all live group pairs, score each with
|
||||
//! average link, and merge the best. It is correct, it is twenty lines, and on
|
||||
//! a real library it does not finish.
|
||||
//!
|
||||
//! The reason is that the scan is inside the loop. Each merge rescans every
|
||||
//! surviving pair — `O(g²)` of them — and each score is recomputed from
|
||||
//! scratch over every cross pair, `O(|A|·|B|)`. With 1,813 faces that is
|
||||
//! roughly 1.6 million pair scores per merge and some 700 merges to do; the
|
||||
//! window simply stops responding, which is what a user reports as "Regroup is
|
||||
//! broken". At 25,000 faces it is not slow, it is impossible.
|
||||
//!
|
||||
//! Three changes, none of which alter the answer:
|
||||
//!
|
||||
//! **Only above-threshold pairs can ever matter.** An average that reaches the
|
||||
//! threshold must have at least one term at or above it, so two groups with no
|
||||
//! qualifying pair between them can never merge — not now and not after any
|
||||
//! sequence of merges, since merging only adds terms. [`crate::neighbours`]
|
||||
//! produces exactly that sparse pair list, and everything below works on it.
|
||||
//! On the library above it is 7,875 pairs rather than 1.6 million.
|
||||
//!
|
||||
//! **Merges cannot cross components.** Groups only ever merge along those
|
||||
//! pairs, so the connected components of that graph are independent problems.
|
||||
//! A library of four hundred people becomes four hundred small agglomerations
|
||||
//! instead of one large one, and the quadratic term is paid per component.
|
||||
//!
|
||||
//! **Average link is additive.** `sum(A ∪ B, C) = sum(A, C) + sum(B, C)`, so a
|
||||
//! merged group's scores follow from the two it came from by addition — the
|
||||
//! Lance-Williams update. Kept as running `(sum, count)` per adjacent pair, a
|
||||
//! score costs one division instead of a nested loop, and a binary heap with
|
||||
//! lazy invalidation replaces the rescan.
|
||||
//!
|
||||
//! The output is unchanged, deliberately and testably so: `the_fast_engine_
|
||||
//! agrees_with_the_reference` runs both over the same population and asserts
|
||||
//! the clusters are identical.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::{BinaryHeap, HashMap, HashSet};
|
||||
|
||||
use crate::calibrate::Calibration;
|
||||
use crate::embedding::EMBEDDING_DIM;
|
||||
use crate::neighbours::{self, Faces};
|
||||
|
||||
/// Probability above which two groups are judged the same person.
|
||||
///
|
||||
@@ -80,79 +119,24 @@ pub fn cluster(faces: &[Candidate], cal: &Calibration, min_probability: f32) ->
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let n = faces.len();
|
||||
let mut groups: Vec<Group> = faces
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, f)| Group {
|
||||
members: vec![i],
|
||||
images: HashSet::from([f.image]),
|
||||
person: f.confirmed_person,
|
||||
alive: true,
|
||||
})
|
||||
.collect();
|
||||
let embeddings: Vec<Vec<f32>> = faces.iter().map(|f| f.embedding.clone()).collect();
|
||||
let crop_px: Vec<f32> = faces.iter().map(|f| f.crop_px).collect();
|
||||
let images: Vec<u64> = faces.iter().map(|f| f.image).collect();
|
||||
let view = Faces {
|
||||
embeddings: &embeddings,
|
||||
crop_px: &crop_px,
|
||||
images: &images,
|
||||
};
|
||||
|
||||
// Pairwise cosine once. n² f32 is the honest cost at library scale — for
|
||||
// 25,000 faces that is the blocked GEMM docs/faces.md §9 describes, and the
|
||||
// caller is expected to shard rather than this function growing an index.
|
||||
let mut cos = vec![0.0_f32; n * n];
|
||||
for i in 0..n {
|
||||
for j in i + 1..n {
|
||||
let c = dot(&faces[i].embedding, &faces[j].embedding);
|
||||
cos[i * n + j] = c;
|
||||
cos[j * n + i] = c;
|
||||
}
|
||||
// Every pair that could ever contribute to a merge. See the module note on
|
||||
// why nothing outside this list can matter.
|
||||
let pairs = neighbours::above_threshold(&view, cal, min_probability);
|
||||
|
||||
let mut engine = Engine::new(faces, cal, min_probability);
|
||||
for component in components(faces.len(), &pairs) {
|
||||
engine.agglomerate(&component, &pairs);
|
||||
}
|
||||
|
||||
loop {
|
||||
let mut best: Option<(f32, usize, usize)> = None;
|
||||
|
||||
for a in 0..n {
|
||||
if !groups[a].alive {
|
||||
continue;
|
||||
}
|
||||
for b in a + 1..n {
|
||||
if !groups[b].alive || !can_link(&groups[a], &groups[b]) {
|
||||
continue;
|
||||
}
|
||||
let p = average_link(&groups[a], &groups[b], faces, &cos, n, cal);
|
||||
if p >= min_probability && best.is_none_or(|(bp, _, _)| p > bp) {
|
||||
best = Some((p, a, b));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let Some((_, a, b)) = best else { break };
|
||||
let taken = std::mem::take(&mut groups[b]);
|
||||
groups[b].alive = false;
|
||||
let ga = &mut groups[a];
|
||||
ga.members.extend(taken.members);
|
||||
ga.images.extend(taken.images);
|
||||
// At most one side carries a person: `can_link` refuses a merge of two
|
||||
// groups anchored to different people, so this cannot silently discard
|
||||
// one of them.
|
||||
ga.person = ga.person.or(taken.person);
|
||||
}
|
||||
|
||||
let mut out: Vec<Cluster> = groups
|
||||
.into_iter()
|
||||
.filter(|g| g.alive)
|
||||
.map(|mut g| {
|
||||
g.members.sort_unstable();
|
||||
Cluster {
|
||||
members: g.members,
|
||||
person: g.person,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
// Largest first: the People view shows the best-evidenced groups at the top.
|
||||
out.sort_by(|x, y| {
|
||||
y.members
|
||||
.len()
|
||||
.cmp(&x.members.len())
|
||||
.then(x.members[0].cmp(&y.members[0]))
|
||||
});
|
||||
out
|
||||
engine.finish()
|
||||
}
|
||||
|
||||
/// Split one person's faces into the groups a raised threshold separates them
|
||||
@@ -178,62 +162,379 @@ pub fn split(faces: &[Candidate], cal: &Calibration, min_probability: f32) -> Ve
|
||||
cluster(&anchorless, cal, min_probability)
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
// ── the merge engine ──────────────────────────────────────────────────────
|
||||
|
||||
/// One live group, identified throughout by the index of its lowest member.
|
||||
#[derive(Debug)]
|
||||
struct Group {
|
||||
/// Ascending, always — [`Engine::cross`] sums in this order, and a stable
|
||||
/// order is what makes the floating-point total reproducible.
|
||||
members: Vec<usize>,
|
||||
images: HashSet<u64>,
|
||||
person: Option<u64>,
|
||||
alive: bool,
|
||||
/// Bumped on every merge, so heap entries naming an older state can be
|
||||
/// recognised and dropped instead of acted on.
|
||||
version: u64,
|
||||
}
|
||||
|
||||
/// Whether two groups are allowed to merge at all, before similarity is asked.
|
||||
fn can_link(a: &Group, b: &Group) -> bool {
|
||||
// Two confirmations of different people. The user has said these are not
|
||||
// the same person, and no similarity overrides that.
|
||||
if let (Some(pa), Some(pb)) = (a.person, b.person) {
|
||||
if pa != pb {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// Co-occurrence: a photograph containing a face from each group means the
|
||||
// two faces are in the same frame, so they are not the same person.
|
||||
a.images.is_disjoint(&b.images)
|
||||
/// Running average-link state for one adjacent pair of groups.
|
||||
///
|
||||
/// `sum` is over **every** cross pair, not only the above-threshold ones —
|
||||
/// average link is an average over all of them, and counting only the
|
||||
/// qualifying pairs would report a similarity no group actually has.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct Link {
|
||||
sum: f64,
|
||||
count: f64,
|
||||
}
|
||||
|
||||
/// Mean calibrated probability over every cross-group pair.
|
||||
fn average_link(
|
||||
a: &Group,
|
||||
b: &Group,
|
||||
faces: &[Candidate],
|
||||
cos: &[f32],
|
||||
n: usize,
|
||||
cal: &Calibration,
|
||||
) -> f32 {
|
||||
let mut sum = 0.0;
|
||||
let mut count = 0.0;
|
||||
for &i in &a.members {
|
||||
for &j in &b.members {
|
||||
let min_crop = faces[i].crop_px.min(faces[j].crop_px);
|
||||
sum += cal.probability(cos[i * n + j], min_crop, 0.0);
|
||||
count += 1.0;
|
||||
impl Link {
|
||||
fn probability(&self) -> f32 {
|
||||
if self.count == 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
(self.sum / self.count) as f32
|
||||
}
|
||||
}
|
||||
if count == 0.0 {
|
||||
0.0
|
||||
}
|
||||
|
||||
/// A candidate merge, waiting in the heap.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct Pending {
|
||||
probability: f32,
|
||||
a: usize,
|
||||
b: usize,
|
||||
/// Group versions when this was pushed. A mismatch on pop means a merge
|
||||
/// has happened since and a fresher entry for this pair is already queued.
|
||||
va: u64,
|
||||
vb: u64,
|
||||
}
|
||||
|
||||
impl PartialEq for Pending {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.cmp(other) == Ordering::Equal
|
||||
}
|
||||
}
|
||||
impl Eq for Pending {}
|
||||
impl PartialOrd for Pending {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
impl Ord for Pending {
|
||||
/// Greatest pops first, so: highest probability, and on a tie the lowest
|
||||
/// index pair. That tiebreak is not cosmetic — it is what the old
|
||||
/// ascending scan did, and it is the whole of the determinism guarantee.
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
self.probability
|
||||
.total_cmp(&other.probability)
|
||||
.then_with(|| other.a.cmp(&self.a))
|
||||
.then_with(|| other.b.cmp(&self.b))
|
||||
}
|
||||
}
|
||||
|
||||
struct Engine<'a> {
|
||||
faces: &'a [Candidate],
|
||||
cal: &'a Calibration,
|
||||
min_probability: f32,
|
||||
groups: Vec<Group>,
|
||||
links: HashMap<(usize, usize), Link>,
|
||||
/// Adjacency, as group ids. Kept alongside `links` so a merge can find
|
||||
/// everything it has to update without scanning the whole map.
|
||||
adjacent: Vec<HashSet<usize>>,
|
||||
}
|
||||
|
||||
impl<'a> Engine<'a> {
|
||||
fn new(faces: &'a [Candidate], cal: &'a Calibration, min_probability: f32) -> Self {
|
||||
let groups = faces
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, f)| Group {
|
||||
members: vec![i],
|
||||
images: HashSet::from([f.image]),
|
||||
person: f.confirmed_person,
|
||||
alive: true,
|
||||
version: 0,
|
||||
})
|
||||
.collect();
|
||||
Self {
|
||||
faces,
|
||||
cal,
|
||||
min_probability,
|
||||
groups,
|
||||
links: HashMap::new(),
|
||||
adjacent: vec![HashSet::new(); faces.len()],
|
||||
}
|
||||
}
|
||||
|
||||
/// Agglomerate one connected component to exhaustion.
|
||||
fn agglomerate(&mut self, component: &[usize], pairs: &[neighbours::Pair]) {
|
||||
if component.len() < 2 {
|
||||
return;
|
||||
}
|
||||
let members: HashSet<usize> = component.iter().copied().collect();
|
||||
|
||||
let mut heap = BinaryHeap::new();
|
||||
for p in pairs.iter().filter(|p| members.contains(&p.i)) {
|
||||
self.links.insert(
|
||||
key(p.i, p.j),
|
||||
Link {
|
||||
sum: p.probability as f64,
|
||||
count: 1.0,
|
||||
},
|
||||
);
|
||||
self.adjacent[p.i].insert(p.j);
|
||||
self.adjacent[p.j].insert(p.i);
|
||||
heap.push(Pending {
|
||||
probability: p.probability,
|
||||
a: p.i.min(p.j),
|
||||
b: p.i.max(p.j),
|
||||
va: 0,
|
||||
vb: 0,
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(top) = heap.pop() {
|
||||
let Pending {
|
||||
probability,
|
||||
a,
|
||||
b,
|
||||
va,
|
||||
vb,
|
||||
} = top;
|
||||
|
||||
// Stale: one side has merged since this was queued, and the
|
||||
// replacement entry is already in the heap.
|
||||
if !self.groups[a].alive
|
||||
|| !self.groups[b].alive
|
||||
|| self.groups[a].version != va
|
||||
|| self.groups[b].version != vb
|
||||
{
|
||||
continue;
|
||||
}
|
||||
// The heap is ordered by probability, so the first entry below the
|
||||
// bar means nothing left in this component can reach it.
|
||||
if probability < self.min_probability {
|
||||
break;
|
||||
}
|
||||
if !self.can_link(a, b) {
|
||||
// Never becomes possible again: images only accumulate and an
|
||||
// anchor is never given up, so drop the pair for good.
|
||||
self.unlink(a, b);
|
||||
continue;
|
||||
}
|
||||
self.merge(a, b, &mut heap);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether two groups are allowed to merge at all, before similarity is
|
||||
/// asked.
|
||||
fn can_link(&self, a: usize, b: usize) -> bool {
|
||||
let (ga, gb) = (&self.groups[a], &self.groups[b]);
|
||||
// Two confirmations of different people. The user has said these are
|
||||
// not the same person, and no similarity overrides that.
|
||||
if let (Some(pa), Some(pb)) = (ga.person, gb.person) {
|
||||
if pa != pb {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// Co-occurrence: a photograph containing a face from each group means
|
||||
// the two faces are in the same frame, so they are not the same person.
|
||||
ga.images.is_disjoint(&gb.images)
|
||||
}
|
||||
|
||||
/// Fold `b` into `a` and re-score everything that touched either.
|
||||
fn merge(&mut self, a: usize, b: usize, heap: &mut BinaryHeap<Pending>) {
|
||||
// Sorted, and deduplicated by the set: `a` and `b` may share
|
||||
// neighbours, and each must be visited once. Sorting is what keeps the
|
||||
// floating-point sums identical from run to run.
|
||||
let mut touched: Vec<usize> = self.adjacent[a]
|
||||
.union(&self.adjacent[b])
|
||||
.copied()
|
||||
.filter(|&c| c != a && c != b && self.groups[c].alive)
|
||||
.collect();
|
||||
touched.sort_unstable();
|
||||
|
||||
// Take the pair sums before the groups change underneath them.
|
||||
let carried: Vec<(usize, Option<Link>, Option<Link>)> = touched
|
||||
.iter()
|
||||
.map(|&c| {
|
||||
(
|
||||
c,
|
||||
self.links.get(&key(a, c)).copied(),
|
||||
self.links.get(&key(b, c)).copied(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// a's own members, before b's are folded in. A missing a-side sum has
|
||||
// to be computed over these and not over the merged list, or b's
|
||||
// contribution would be counted twice.
|
||||
let a_members = self.groups[a].members.clone();
|
||||
|
||||
// Absorb b into a.
|
||||
let taken = std::mem::replace(
|
||||
&mut self.groups[b],
|
||||
Group {
|
||||
members: Vec::new(),
|
||||
images: HashSet::new(),
|
||||
person: None,
|
||||
alive: false,
|
||||
version: 0,
|
||||
},
|
||||
);
|
||||
{
|
||||
let ga = &mut self.groups[a];
|
||||
ga.members.extend(taken.members.iter().copied());
|
||||
ga.members.sort_unstable();
|
||||
ga.images.extend(taken.images.iter().copied());
|
||||
// At most one side carries a person: `can_link` refuses a merge of
|
||||
// two groups anchored to different people, so this cannot silently
|
||||
// discard one of them.
|
||||
ga.person = ga.person.or(taken.person);
|
||||
ga.version += 1;
|
||||
}
|
||||
|
||||
// b's own links are gone with it.
|
||||
for c in self.adjacent[b].clone() {
|
||||
self.links.remove(&key(b, c));
|
||||
self.adjacent[c].remove(&b);
|
||||
}
|
||||
self.adjacent[b].clear();
|
||||
self.links.remove(&key(a, b));
|
||||
self.adjacent[a].remove(&b);
|
||||
|
||||
for (c, from_a, from_b) in carried {
|
||||
// Dropping a pair the constraints now forbid saves computing a
|
||||
// score for a merge that can never happen — which for a newly
|
||||
// adjacent side is a real cost, not a bookkeeping one.
|
||||
if !self.can_link(a, c) {
|
||||
self.unlink(a, c);
|
||||
continue;
|
||||
}
|
||||
// A side with no stored link was not adjacent before, so its cross
|
||||
// pairs were all below threshold and were never summed. They still
|
||||
// belong in the average, so they are computed now — once, after
|
||||
// which the additive update carries them forward.
|
||||
let from_a = from_a.unwrap_or_else(|| self.cross(&a_members, c));
|
||||
let from_b = from_b.unwrap_or_else(|| self.cross(&taken.members, c));
|
||||
let merged = Link {
|
||||
sum: from_a.sum + from_b.sum,
|
||||
count: from_a.count + from_b.count,
|
||||
};
|
||||
self.links.insert(key(a, c), merged);
|
||||
self.adjacent[a].insert(c);
|
||||
self.adjacent[c].insert(a);
|
||||
|
||||
heap.push(Pending {
|
||||
probability: merged.probability(),
|
||||
a: a.min(c),
|
||||
b: a.max(c),
|
||||
va: self.groups[a.min(c)].version,
|
||||
vb: self.groups[a.max(c)].version,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Exact `(sum, count)` over every cross pair between a member list and a
|
||||
/// group.
|
||||
///
|
||||
/// The one place a dot product is still computed during agglomeration, and
|
||||
/// it happens only when two groups become adjacent through a third — at
|
||||
/// which point their sub-threshold pairs, never summed because they were
|
||||
/// never interesting, have to be accounted for.
|
||||
fn cross(&self, members: &[usize], group: usize) -> Link {
|
||||
let mut sum = 0.0_f64;
|
||||
let mut count = 0.0_f64;
|
||||
for &i in members {
|
||||
for &j in &self.groups[group].members {
|
||||
let cos = neighbours::dot(&self.faces[i].embedding, &self.faces[j].embedding);
|
||||
let min_crop = self.faces[i].crop_px.min(self.faces[j].crop_px);
|
||||
sum += self.cal.probability(cos, min_crop, 0.0) as f64;
|
||||
count += 1.0;
|
||||
}
|
||||
}
|
||||
Link { sum, count }
|
||||
}
|
||||
|
||||
fn unlink(&mut self, a: usize, b: usize) {
|
||||
self.links.remove(&key(a, b));
|
||||
self.adjacent[a].remove(&b);
|
||||
self.adjacent[b].remove(&a);
|
||||
}
|
||||
|
||||
fn finish(self) -> Vec<Cluster> {
|
||||
let mut out: Vec<Cluster> = self
|
||||
.groups
|
||||
.into_iter()
|
||||
.filter(|g| g.alive)
|
||||
.map(|g| Cluster {
|
||||
members: g.members,
|
||||
person: g.person,
|
||||
})
|
||||
.collect();
|
||||
// Largest first: the People view shows the best-evidenced groups at the
|
||||
// top.
|
||||
out.sort_by(|x, y| {
|
||||
y.members
|
||||
.len()
|
||||
.cmp(&x.members.len())
|
||||
.then(x.members[0].cmp(&y.members[0]))
|
||||
});
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
fn key(a: usize, b: usize) -> (usize, usize) {
|
||||
if a < b {
|
||||
(a, b)
|
||||
} else {
|
||||
sum / count
|
||||
(b, a)
|
||||
}
|
||||
}
|
||||
|
||||
fn dot(a: &[f32], b: &[f32]) -> f32 {
|
||||
debug_assert_eq!(a.len(), EMBEDDING_DIM);
|
||||
debug_assert_eq!(b.len(), EMBEDDING_DIM);
|
||||
a.iter().zip(b).map(|(x, y)| x * y).sum()
|
||||
}
|
||||
/// Connected components of the above-threshold graph.
|
||||
///
|
||||
/// Faces in different components can never end up in one group, so each is a
|
||||
/// separate and much smaller agglomeration. Returned with the members of each
|
||||
/// component ascending, and the components themselves in order of their lowest
|
||||
/// member — the determinism the merge order inherits.
|
||||
fn components(n: usize, pairs: &[neighbours::Pair]) -> Vec<Vec<usize>> {
|
||||
let mut parent: Vec<usize> = (0..n).collect();
|
||||
|
||||
fn find(parent: &mut [usize], mut x: usize) -> usize {
|
||||
while parent[x] != x {
|
||||
// Path halving: keeps the tree flat without a second pass.
|
||||
parent[x] = parent[parent[x]];
|
||||
x = parent[x];
|
||||
}
|
||||
x
|
||||
}
|
||||
|
||||
for p in pairs {
|
||||
let (ra, rb) = (find(&mut parent, p.i), find(&mut parent, p.j));
|
||||
if ra != rb {
|
||||
// Lowest root wins, so the representative of a component is
|
||||
// reproducible rather than an artefact of union order.
|
||||
let (lo, hi) = if ra < rb { (ra, rb) } else { (rb, ra) };
|
||||
parent[hi] = lo;
|
||||
}
|
||||
}
|
||||
|
||||
let mut by_root: HashMap<usize, Vec<usize>> = HashMap::new();
|
||||
for i in 0..n {
|
||||
let r = find(&mut parent, i);
|
||||
by_root.entry(r).or_default().push(i);
|
||||
}
|
||||
let mut out: Vec<Vec<usize>> = by_root.into_values().filter(|c| c.len() > 1).collect();
|
||||
out.sort_unstable_by_key(|c| c[0]);
|
||||
out
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::embedding::EMBEDDING_DIM;
|
||||
|
||||
/// An embedding a known cosine away from a base direction, built by mixing
|
||||
/// two orthogonal unit vectors. Lets a test state "these two faces are 0.7
|
||||
@@ -441,4 +742,221 @@ mod tests {
|
||||
let small = sized.probability(0.5, 40.0, 0.0);
|
||||
assert!(big > small, "big {big} should beat small {small}");
|
||||
}
|
||||
|
||||
// ── the fast engine against the obvious one ───────────────────────────
|
||||
|
||||
/// The original implementation, kept as the oracle.
|
||||
///
|
||||
/// Deliberately the naive version this module replaced: rescan every live
|
||||
/// pair, score it from scratch over all cross pairs, merge the best,
|
||||
/// repeat. It is the definition of the answer, and the only thing the
|
||||
/// rewrite was allowed to change is how long it takes to get there.
|
||||
fn reference(faces: &[Candidate], cal: &Calibration, min_probability: f32) -> Vec<Cluster> {
|
||||
#[derive(Clone)]
|
||||
struct G {
|
||||
members: Vec<usize>,
|
||||
images: HashSet<u64>,
|
||||
person: Option<u64>,
|
||||
alive: bool,
|
||||
}
|
||||
if faces.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let n = faces.len();
|
||||
let mut groups: Vec<G> = faces
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, f)| G {
|
||||
members: vec![i],
|
||||
images: HashSet::from([f.image]),
|
||||
person: f.confirmed_person,
|
||||
alive: true,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut cos = vec![0.0_f32; n * n];
|
||||
for i in 0..n {
|
||||
for j in i + 1..n {
|
||||
let c = neighbours::dot(&faces[i].embedding, &faces[j].embedding);
|
||||
cos[i * n + j] = c;
|
||||
cos[j * n + i] = c;
|
||||
}
|
||||
}
|
||||
let linkable = |a: &G, b: &G| {
|
||||
if let (Some(pa), Some(pb)) = (a.person, b.person) {
|
||||
if pa != pb {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
a.images.is_disjoint(&b.images)
|
||||
};
|
||||
let average = |a: &G, b: &G| {
|
||||
let mut sum = 0.0_f32;
|
||||
let mut count = 0.0_f32;
|
||||
for &i in &a.members {
|
||||
for &j in &b.members {
|
||||
let min_crop = faces[i].crop_px.min(faces[j].crop_px);
|
||||
sum += cal.probability(cos[i * n + j], min_crop, 0.0);
|
||||
count += 1.0;
|
||||
}
|
||||
}
|
||||
if count == 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
sum / count
|
||||
}
|
||||
};
|
||||
|
||||
loop {
|
||||
let mut best: Option<(f32, usize, usize)> = None;
|
||||
for a in 0..n {
|
||||
if !groups[a].alive {
|
||||
continue;
|
||||
}
|
||||
for b in a + 1..n {
|
||||
if !groups[b].alive || !linkable(&groups[a], &groups[b]) {
|
||||
continue;
|
||||
}
|
||||
let p = average(&groups[a], &groups[b]);
|
||||
if p >= min_probability && best.is_none_or(|(bp, _, _)| p > bp) {
|
||||
best = Some((p, a, b));
|
||||
}
|
||||
}
|
||||
}
|
||||
let Some((_, a, b)) = best else { break };
|
||||
let taken = groups[b].clone();
|
||||
groups[b].alive = false;
|
||||
groups[a].members.extend(taken.members);
|
||||
groups[a].images.extend(taken.images);
|
||||
groups[a].person = groups[a].person.or(taken.person);
|
||||
}
|
||||
|
||||
let mut out: Vec<Cluster> = groups
|
||||
.into_iter()
|
||||
.filter(|g| g.alive)
|
||||
.map(|mut g| {
|
||||
g.members.sort_unstable();
|
||||
Cluster {
|
||||
members: g.members,
|
||||
person: g.person,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
out.sort_by(|x, y| {
|
||||
y.members
|
||||
.len()
|
||||
.cmp(&x.members.len())
|
||||
.then(x.members[0].cmp(&y.members[0]))
|
||||
});
|
||||
out
|
||||
}
|
||||
|
||||
/// `people` identities of `per` faces, each face in its own photograph,
|
||||
/// spread either side of the threshold so the population has genuine
|
||||
/// near-misses rather than obvious answers.
|
||||
fn population(people: usize, per: usize) -> Vec<Candidate> {
|
||||
let mut out = Vec::new();
|
||||
let mut image = 0u64;
|
||||
for p in 0..people {
|
||||
for m in 0..per {
|
||||
// Walks down through the merge boundary as m grows, so some
|
||||
// members join their group and some do not.
|
||||
let cosine = 1.0 - (m as f32) * 0.035;
|
||||
out.push(Candidate {
|
||||
face: out.len() as u64,
|
||||
image,
|
||||
embedding: at_cosine(p, cosine),
|
||||
crop_px: 60.0 + ((out.len() % 11) as f32) * 25.0,
|
||||
confirmed_person: None,
|
||||
});
|
||||
image += 1;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// The point of the rewrite: same clusters, less work. A disagreement here
|
||||
/// is the rewrite being wrong, not the reference being slow.
|
||||
#[test]
|
||||
fn the_fast_engine_agrees_with_the_reference() {
|
||||
let faces = population(40, 6);
|
||||
assert_eq!(
|
||||
cluster(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
reference(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
);
|
||||
}
|
||||
|
||||
/// The two structural constraints are the ones a sparse graph could
|
||||
/// plausibly break, so they get their own comparison with anchors and
|
||||
/// co-occurrence in play.
|
||||
#[test]
|
||||
fn the_fast_engine_agrees_with_the_reference_under_constraints() {
|
||||
let mut faces = population(30, 6);
|
||||
// Some faces share a photograph, so cannot-link has to propagate
|
||||
// through groups that formed for other reasons.
|
||||
for i in (0..faces.len()).step_by(7) {
|
||||
faces[i].image = 900 + (i as u64 % 4);
|
||||
}
|
||||
// And some carry confirmations, including two of different people that
|
||||
// must never be brought together.
|
||||
for (n, i) in (0..faces.len()).step_by(11).enumerate() {
|
||||
faces[i].confirmed_person = Some(1 + (n as u64 % 3));
|
||||
}
|
||||
assert_eq!(
|
||||
cluster(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
reference(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
);
|
||||
}
|
||||
|
||||
/// The size term makes the merge boundary depend on the pair, which is the
|
||||
/// case the sparse pre-filter has to be built carefully to preserve.
|
||||
#[test]
|
||||
fn the_fast_engine_agrees_with_the_reference_with_a_size_term() {
|
||||
let faces = population(30, 6);
|
||||
let sized = Calibration {
|
||||
w_size: 0.4,
|
||||
b: -30.0 * 0.35 - 0.4 * 7.0,
|
||||
..cal()
|
||||
};
|
||||
assert_eq!(
|
||||
cluster(&faces, &sized, DEFAULT_MERGE_PROBABILITY),
|
||||
reference(&faces, &sized, DEFAULT_MERGE_PROBABILITY),
|
||||
);
|
||||
}
|
||||
|
||||
/// Determinism has to hold at a size where the indexed neighbour search is
|
||||
/// in play, not just on the handful of faces the small cases use.
|
||||
#[test]
|
||||
fn clustering_is_deterministic_at_scale() {
|
||||
let faces = population(200, 6);
|
||||
assert!(
|
||||
faces.len() > 1024,
|
||||
"population is below the indexing cutoff"
|
||||
);
|
||||
assert_eq!(
|
||||
cluster(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
cluster(&faces, &cal(), DEFAULT_MERGE_PROBABILITY),
|
||||
);
|
||||
}
|
||||
|
||||
/// A face that matches nobody is left alone rather than being swept into
|
||||
/// the nearest group, and costs nothing to establish — it is in no
|
||||
/// component at all.
|
||||
#[test]
|
||||
fn a_face_matching_nothing_stays_on_its_own() {
|
||||
let mut faces = population(5, 4);
|
||||
faces.push(Candidate {
|
||||
face: 999,
|
||||
image: 5_000,
|
||||
embedding: at_cosine(200, 1.0),
|
||||
crop_px: 150.0,
|
||||
confirmed_person: None,
|
||||
});
|
||||
let out = cluster(&faces, &cal(), DEFAULT_MERGE_PROBABILITY);
|
||||
let last = faces.len() - 1;
|
||||
assert!(
|
||||
out.iter().any(|c| c.members == vec![last]),
|
||||
"the outlier was absorbed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user