videre_core/
face_db.rs

1use half::f16;
2use rusqlite::Connection;
3use std::collections::HashMap;
4
5pub struct FaceRow {
6    pub hash: String,
7    pub bbox: String,
8    pub landmark: Option<String>,
9    pub embedding: Vec<u8>, // 512 f16 values as little-endian bytes (1024 bytes)
10    pub cluster_id: Option<i64>,
11    pub person_label: Option<String>,
12    pub confirmed: i64,
13    pub is_primary: i64,
14}
15
16pub fn create_faces_table(conn: &Connection) -> rusqlite::Result<()> {
17    conn.execute_batch(
18        "CREATE TABLE IF NOT EXISTS faces (
19            id            INTEGER PRIMARY KEY,
20            hash          TEXT NOT NULL,
21            bbox          TEXT NOT NULL,
22            landmark      TEXT,
23            embedding     BLOB NOT NULL,
24            cluster_id    INTEGER,
25            person_label  TEXT,
26            confirmed     INTEGER DEFAULT 0,
27            is_primary    INTEGER DEFAULT 0
28        );",
29    )?;
30    // Migration for existing tables without is_primary column; ignored if already exists.
31    let _ = conn.execute_batch("ALTER TABLE faces ADD COLUMN is_primary INTEGER DEFAULT 0");
32    // Records every hash whose faces have been scanned, INCLUDING images where
33    // zero faces were detected (which leave no `faces` row). This is what makes
34    // `videre faces` resumable: the skip set is "already scanned", not merely
35    // "has a face", so a no-face image is never re-detected on a later run.
36    conn.execute_batch(
37        "CREATE TABLE IF NOT EXISTS faces_scanned (
38            hash        TEXT PRIMARY KEY,
39            scanned_at  TEXT DEFAULT (datetime('now'))
40        );",
41    )?;
42    Ok(())
43}
44
45/// Marks a hash as face-scanned (idempotent). Call after detection runs for a
46/// hash regardless of whether any faces were found.
47pub fn mark_scanned(conn: &Connection, hash: &str) -> rusqlite::Result<()> {
48    conn.execute(
49        "INSERT OR IGNORE INTO faces_scanned (hash) VALUES (?1)",
50        rusqlite::params![hash],
51    )?;
52    Ok(())
53}
54
55/// Every hash recorded as face-scanned.
56pub fn scanned_hashes(conn: &Connection) -> rusqlite::Result<Vec<String>> {
57    let mut stmt = conn.prepare("SELECT hash FROM faces_scanned")?;
58    let rows = stmt.query_map([], |r| r.get(0))?;
59    rows.collect()
60}
61
62/// From `(path, hash)` pairs, drop hashes in `skip`, keep one representative
63/// path per remaining hash (first seen), preserving input order, and cap the
64/// result at `limit` distinct hashes (`None` = no cap). Used to build the work
65/// list for a resumable, optionally partial face-detection pass.
66pub fn select_unscanned(
67    all: &[(String, String)],
68    skip: &std::collections::HashSet<String>,
69    limit: Option<usize>,
70) -> Vec<(String, String)> {
71    let mut seen = std::collections::HashSet::new();
72    let mut out = Vec::new();
73    for (path, hash) in all {
74        if skip.contains(hash) || !seen.insert(hash.clone()) {
75            continue;
76        }
77        out.push((path.clone(), hash.clone()));
78        if let Some(n) = limit {
79            if out.len() >= n {
80                break;
81            }
82        }
83    }
84    out
85}
86
87pub fn replace_faces_for_hash(
88    conn: &Connection,
89    hash: &str,
90    faces: &[FaceRow],
91) -> rusqlite::Result<()> {
92    conn.execute_batch("BEGIN")?;
93    let result = (|| -> rusqlite::Result<()> {
94        conn.execute("DELETE FROM faces WHERE hash = ?1", rusqlite::params![hash])?;
95        for face in faces {
96            conn.execute(
97                "INSERT INTO faces (hash, bbox, landmark, embedding, cluster_id, person_label, confirmed, is_primary)
98                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
99                rusqlite::params![
100                    face.hash, face.bbox, face.landmark, face.embedding,
101                    face.cluster_id, face.person_label, face.confirmed, face.is_primary
102                ],
103            )?;
104        }
105        Ok(())
106    })();
107    match result {
108        Ok(()) => {
109            conn.execute_batch("COMMIT")?;
110            Ok(())
111        }
112        Err(e) => {
113            let _ = conn.execute_batch("ROLLBACK");
114            Err(e)
115        }
116    }
117}
118
119pub fn load_face_embeddings(conn: &Connection) -> rusqlite::Result<Vec<(i64, Vec<f32>)>> {
120    let mut stmt = conn.prepare("SELECT id, embedding FROM faces")?;
121    let rows = stmt.query_map([], |row| {
122        let id: i64 = row.get(0)?;
123        let blob: Vec<u8> = row.get(1)?;
124        Ok((id, blob))
125    })?;
126    let mut out = Vec::new();
127    for row in rows {
128        let (id, blob) = row?;
129        let emb: Vec<f32> = blob
130            .chunks_exact(2)
131            .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
132            .collect();
133        out.push((id, emb));
134    }
135    Ok(out)
136}
137
138/// Like [`load_face_embeddings`] but also returns each face's smaller bbox
139/// side in pixels (the shorter of width/height), parsed from the `"x,y,w,h"`
140/// bbox string. Used as a quality signal: very small face crops embed into
141/// near-degenerate ArcFace vectors that cluster together regardless of
142/// identity, so callers gate them out of clustering. A bbox that fails to
143/// parse yields a min-side of 0.0 (treated as lowest quality).
144pub fn load_faces_for_clustering(conn: &Connection) -> rusqlite::Result<Vec<(i64, Vec<f32>, f32)>> {
145    let mut stmt = conn.prepare("SELECT id, embedding, bbox FROM faces")?;
146    let rows = stmt.query_map([], |row| {
147        let id: i64 = row.get(0)?;
148        let blob: Vec<u8> = row.get(1)?;
149        let bbox: String = row.get(2)?;
150        Ok((id, blob, bbox))
151    })?;
152    let mut out = Vec::new();
153    for row in rows {
154        let (id, blob, bbox) = row?;
155        let emb: Vec<f32> = blob
156            .chunks_exact(2)
157            .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
158            .collect();
159        out.push((id, emb, bbox_min_side(&bbox)));
160    }
161    Ok(out)
162}
163
164/// Smaller side (min of width, height) of a `"x,y,w,h"` bbox string, or 0.0 if
165/// it does not parse into at least four numeric fields.
166fn bbox_min_side(bbox: &str) -> f32 {
167    let nums: Vec<f32> = bbox
168        .split(',')
169        .filter_map(|s| s.trim().parse().ok())
170        .collect();
171    if nums.len() >= 4 {
172        nums[2].min(nums[3])
173    } else {
174        0.0
175    }
176}
177
178pub fn update_cluster_assignments(
179    conn: &Connection,
180    assignments: &[(i64, Option<i64>)],
181) -> rusqlite::Result<()> {
182    for (face_id, cluster_id) in assignments {
183        conn.execute(
184            "UPDATE faces SET cluster_id = ?1 WHERE id = ?2",
185            rusqlite::params![cluster_id, face_id],
186        )?;
187    }
188    Ok(())
189}
190
191pub fn hashes_with_faces(conn: &Connection) -> rusqlite::Result<Vec<String>> {
192    let mut stmt = conn.prepare("SELECT DISTINCT hash FROM faces ORDER BY hash")?;
193    let rows = stmt.query_map([], |r| r.get(0))?;
194    rows.collect()
195}
196
197/// (face_id, person_label, bbox) for one labeled face.
198pub type LabeledFace = (i64, String, String);
199
200/// Maps a file hash to every labeled face on it, as returned by
201/// `labeled_faces_by_hash`.
202pub type LabeledFacesByHash = HashMap<String, Vec<LabeledFace>>;
203
204/// Returns, for every hash that has at least one confirmed+labeled face, the
205/// list of (face_id, person_label, bbox) for that hash. One batched query
206/// covering every hash, not one query per file, safe to call once per
207/// report generation without N+1 overhead.
208pub fn labeled_faces_by_hash(conn: &Connection) -> rusqlite::Result<LabeledFacesByHash> {
209    let mut stmt = conn.prepare(
210        "SELECT hash, id, bbox, person_label FROM faces \
211         WHERE confirmed = 1 AND person_label IS NOT NULL \
212         ORDER BY hash, id",
213    )?;
214    let rows = stmt.query_map([], |r| {
215        Ok((
216            r.get::<_, String>(0)?,
217            r.get::<_, i64>(1)?,
218            r.get::<_, String>(2)?,
219            r.get::<_, String>(3)?,
220        ))
221    })?;
222    let mut map: LabeledFacesByHash = HashMap::new();
223    for row in rows {
224        let (hash, id, bbox, label) = row?;
225        map.entry(hash).or_default().push((id, label, bbox));
226    }
227    Ok(map)
228}
229
230#[cfg(test)]
231fn make_embedding(vals: &[f32]) -> Vec<u8> {
232    vals.iter()
233        .flat_map(|&v| f16::from_f32(v).to_le_bytes())
234        .collect()
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    fn open() -> Connection {
242        let conn = Connection::open_in_memory().unwrap();
243        create_faces_table(&conn).unwrap();
244        conn
245    }
246
247    #[test]
248    fn create_table_idempotent() {
249        let conn = open();
250        create_faces_table(&conn).unwrap();
251    }
252
253    #[test]
254    fn insert_and_load_embedding() {
255        let conn = open();
256        let emb = make_embedding(&vec![0.5f32; 512]);
257        replace_faces_for_hash(
258            &conn,
259            "habc",
260            &[FaceRow {
261                hash: "habc".into(),
262                bbox: "0,0,50,50".into(),
263                landmark: None,
264                embedding: emb,
265                cluster_id: None,
266                person_label: None,
267                confirmed: 0,
268                is_primary: 0,
269            }],
270        )
271        .unwrap();
272        let rows = load_face_embeddings(&conn).unwrap();
273        assert_eq!(rows.len(), 1);
274        let (id, emb_f32) = &rows[0];
275        assert!(*id > 0);
276        assert_eq!(emb_f32.len(), 512);
277        assert!((emb_f32[0] - 0.5).abs() < 0.01);
278    }
279
280    #[test]
281    fn replace_removes_old_rows_for_same_hash() {
282        let conn = open();
283        let emb = make_embedding(&vec![0.0f32; 512]);
284        replace_faces_for_hash(
285            &conn,
286            "h1",
287            &[
288                FaceRow {
289                    hash: "h1".into(),
290                    bbox: "0,0,10,10".into(),
291                    landmark: None,
292                    embedding: emb.clone(),
293                    cluster_id: None,
294                    person_label: None,
295                    confirmed: 0,
296                    is_primary: 0,
297                },
298                FaceRow {
299                    hash: "h1".into(),
300                    bbox: "20,0,10,10".into(),
301                    landmark: None,
302                    embedding: emb.clone(),
303                    cluster_id: None,
304                    person_label: None,
305                    confirmed: 0,
306                    is_primary: 0,
307                },
308            ],
309        )
310        .unwrap();
311        replace_faces_for_hash(
312            &conn,
313            "h1",
314            &[FaceRow {
315                hash: "h1".into(),
316                bbox: "99,0,10,10".into(),
317                landmark: None,
318                embedding: emb,
319                cluster_id: None,
320                person_label: None,
321                confirmed: 0,
322                is_primary: 0,
323            }],
324        )
325        .unwrap();
326        let rows = load_face_embeddings(&conn).unwrap();
327        assert_eq!(rows.len(), 1);
328    }
329
330    #[test]
331    fn update_cluster_assignments_works() {
332        let conn = open();
333        let emb = make_embedding(&vec![0.0f32; 512]);
334        replace_faces_for_hash(
335            &conn,
336            "h1",
337            &[FaceRow {
338                hash: "h1".into(),
339                bbox: "0,0,10,10".into(),
340                landmark: None,
341                embedding: emb,
342                cluster_id: None,
343                person_label: None,
344                confirmed: 0,
345                is_primary: 0,
346            }],
347        )
348        .unwrap();
349        let rows = load_face_embeddings(&conn).unwrap();
350        let id = rows[0].0;
351        update_cluster_assignments(&conn, &[(id, Some(3))]).unwrap();
352        let n: i64 = conn
353            .query_row("SELECT cluster_id FROM faces WHERE id=?1", [id], |r| {
354                r.get(0)
355            })
356            .unwrap();
357        assert_eq!(n, 3);
358    }
359
360    #[test]
361    fn load_faces_for_clustering_returns_bbox_min_side() {
362        let conn = open();
363        let emb = make_embedding(&vec![0.25f32; 512]);
364        // bbox "x,y,w,h": min side is min(w,h).
365        replace_faces_for_hash(
366            &conn,
367            "h1",
368            &[
369                FaceRow {
370                    hash: "h1".into(),
371                    bbox: "10,10,200,300".into(),
372                    landmark: None,
373                    embedding: emb.clone(),
374                    cluster_id: None,
375                    person_label: None,
376                    confirmed: 0,
377                    is_primary: 0,
378                },
379                FaceRow {
380                    hash: "h1".into(),
381                    bbox: "0,0,40,25".into(),
382                    landmark: None,
383                    embedding: emb,
384                    cluster_id: None,
385                    person_label: None,
386                    confirmed: 0,
387                    is_primary: 0,
388                },
389            ],
390        )
391        .unwrap();
392        let mut rows = load_faces_for_clustering(&conn).unwrap();
393        rows.sort_by(|a, b| b.2.total_cmp(&a.2));
394        assert_eq!(rows[0].2, 200.0, "min side of 200x300 bbox");
395        assert_eq!(rows[1].2, 25.0, "min side of 40x25 bbox");
396        assert_eq!(rows[0].1.len(), 512, "embedding still decoded");
397    }
398
399    #[test]
400    fn mark_scanned_records_hash_even_with_zero_faces() {
401        let conn = open();
402        // A hash processed with no detected faces leaves no `faces` row, but
403        // must still be recorded as scanned so it is not re-processed.
404        mark_scanned(&conn, "noface").unwrap();
405        assert_eq!(scanned_hashes(&conn).unwrap(), vec!["noface".to_string()]);
406        // hashes_with_faces stays empty, the marker is independent of faces.
407        assert!(hashes_with_faces(&conn).unwrap().is_empty());
408    }
409
410    #[test]
411    fn mark_scanned_is_idempotent() {
412        let conn = open();
413        mark_scanned(&conn, "h").unwrap();
414        mark_scanned(&conn, "h").unwrap();
415        assert_eq!(scanned_hashes(&conn).unwrap().len(), 1);
416    }
417
418    #[test]
419    fn select_unscanned_skips_dedups_and_limits() {
420        // Two paths share hash "a"; "b" is skipped; "c","d","e" remain.
421        let all = vec![
422            ("/1.jpg".to_string(), "a".to_string()),
423            ("/1copy.jpg".to_string(), "a".to_string()),
424            ("/2.jpg".to_string(), "b".to_string()),
425            ("/3.jpg".to_string(), "c".to_string()),
426            ("/4.jpg".to_string(), "d".to_string()),
427            ("/5.jpg".to_string(), "e".to_string()),
428        ];
429        let skip: std::collections::HashSet<String> = ["b".to_string()].into_iter().collect();
430        // No limit: one path per unscanned hash (a,c,d,e), b excluded.
431        let out = select_unscanned(&all, &skip, None);
432        assert_eq!(
433            out.iter().map(|(_, h)| h.clone()).collect::<Vec<_>>(),
434            vec!["a", "c", "d", "e"]
435        );
436        // Limit 2: first two unscanned hashes only.
437        let out2 = select_unscanned(&all, &skip, Some(2));
438        assert_eq!(
439            out2.iter().map(|(_, h)| h.clone()).collect::<Vec<_>>(),
440            vec!["a", "c"]
441        );
442    }
443
444    #[test]
445    fn hashes_with_faces_returns_inserted_hash() {
446        let conn = open();
447        let emb = make_embedding(&vec![0.0f32; 512]);
448        replace_faces_for_hash(
449            &conn,
450            "myhash",
451            &[FaceRow {
452                hash: "myhash".into(),
453                bbox: "0,0,10,10".into(),
454                landmark: None,
455                embedding: emb,
456                cluster_id: None,
457                person_label: None,
458                confirmed: 0,
459                is_primary: 0,
460            }],
461        )
462        .unwrap();
463        let hashes = hashes_with_faces(&conn).unwrap();
464        assert_eq!(hashes, vec!["myhash"]);
465    }
466
467    #[test]
468    fn labeled_faces_by_hash_returns_only_confirmed_labeled() {
469        let conn = Connection::open_in_memory().unwrap();
470        create_faces_table(&conn).unwrap();
471        conn.execute_batch(
472            "INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
473             VALUES ('h1', '0,0,10,10', X'0000', 'Alice', 1); \
474             INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
475             VALUES ('h1', '20,20,10,10', X'0000', NULL, 0); \
476             INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
477             VALUES ('h2', '0,0,10,10', X'0000', 'Bob', 1);",
478        )
479        .unwrap();
480
481        let map = labeled_faces_by_hash(&conn).unwrap();
482        assert_eq!(map.len(), 2, "expected two hashes with labeled faces");
483        let h1 = &map["h1"];
484        assert_eq!(h1.len(), 1, "unconfirmed/unlabeled face must be excluded");
485        assert_eq!(h1[0].1, "Alice");
486        assert_eq!(map["h2"][0].1, "Bob");
487    }
488}
videre_core/face_db.rs

videre_core/
face_db.rs
