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("INSERT OR IGNORE INTO faces_scanned (hash) VALUES (?1)", rusqlite::params![hash])?;
49    Ok(())
50}
51
52/// Every hash recorded as face-scanned.
53pub fn scanned_hashes(conn: &Connection) -> rusqlite::Result<Vec<String>> {
54    let mut stmt = conn.prepare("SELECT hash FROM faces_scanned")?;
55    let rows = stmt.query_map([], |r| r.get(0))?;
56    rows.collect()
57}
58
59/// From `(path, hash)` pairs, drop hashes in `skip`, keep one representative
60/// path per remaining hash (first seen), preserving input order, and cap the
61/// result at `limit` distinct hashes (`None` = no cap). Used to build the work
62/// list for a resumable, optionally partial face-detection pass.
63pub fn select_unscanned(
64    all: &[(String, String)],
65    skip: &std::collections::HashSet<String>,
66    limit: Option<usize>,
67) -> Vec<(String, String)> {
68    let mut seen = std::collections::HashSet::new();
69    let mut out = Vec::new();
70    for (path, hash) in all {
71        if skip.contains(hash) || !seen.insert(hash.clone()) {
72            continue;
73        }
74        out.push((path.clone(), hash.clone()));
75        if let Some(n) = limit {
76            if out.len() >= n {
77                break;
78            }
79        }
80    }
81    out
82}
83
84pub fn replace_faces_for_hash(conn: &Connection, hash: &str, faces: &[FaceRow]) -> rusqlite::Result<()> {
85    conn.execute_batch("BEGIN")?;
86    let result = (|| -> rusqlite::Result<()> {
87        conn.execute("DELETE FROM faces WHERE hash = ?1", rusqlite::params![hash])?;
88        for face in faces {
89            conn.execute(
90                "INSERT INTO faces (hash, bbox, landmark, embedding, cluster_id, person_label, confirmed, is_primary)
91                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
92                rusqlite::params![
93                    face.hash, face.bbox, face.landmark, face.embedding,
94                    face.cluster_id, face.person_label, face.confirmed, face.is_primary
95                ],
96            )?;
97        }
98        Ok(())
99    })();
100    match result {
101        Ok(()) => { conn.execute_batch("COMMIT")?; Ok(()) }
102        Err(e) => { let _ = conn.execute_batch("ROLLBACK"); Err(e) }
103    }
104}
105
106pub fn load_face_embeddings(conn: &Connection) -> rusqlite::Result<Vec<(i64, Vec<f32>)>> {
107    let mut stmt = conn.prepare("SELECT id, embedding FROM faces")?;
108    let rows = stmt.query_map([], |row| {
109        let id: i64 = row.get(0)?;
110        let blob: Vec<u8> = row.get(1)?;
111        Ok((id, blob))
112    })?;
113    let mut out = Vec::new();
114    for row in rows {
115        let (id, blob) = row?;
116        let emb: Vec<f32> = blob
117            .chunks_exact(2)
118            .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
119            .collect();
120        out.push((id, emb));
121    }
122    Ok(out)
123}
124
125/// Like [`load_face_embeddings`] but also returns each face's smaller bbox
126/// side in pixels (the shorter of width/height), parsed from the `"x,y,w,h"`
127/// bbox string. Used as a quality signal: very small face crops embed into
128/// near-degenerate ArcFace vectors that cluster together regardless of
129/// identity, so callers gate them out of clustering. A bbox that fails to
130/// parse yields a min-side of 0.0 (treated as lowest quality).
131pub fn load_faces_for_clustering(conn: &Connection) -> rusqlite::Result<Vec<(i64, Vec<f32>, f32)>> {
132    let mut stmt = conn.prepare("SELECT id, embedding, bbox FROM faces")?;
133    let rows = stmt.query_map([], |row| {
134        let id: i64 = row.get(0)?;
135        let blob: Vec<u8> = row.get(1)?;
136        let bbox: String = row.get(2)?;
137        Ok((id, blob, bbox))
138    })?;
139    let mut out = Vec::new();
140    for row in rows {
141        let (id, blob, bbox) = row?;
142        let emb: Vec<f32> = blob
143            .chunks_exact(2)
144            .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
145            .collect();
146        out.push((id, emb, bbox_min_side(&bbox)));
147    }
148    Ok(out)
149}
150
151/// Smaller side (min of width, height) of a `"x,y,w,h"` bbox string, or 0.0 if
152/// it does not parse into at least four numeric fields.
153fn bbox_min_side(bbox: &str) -> f32 {
154    let nums: Vec<f32> = bbox.split(',').filter_map(|s| s.trim().parse().ok()).collect();
155    if nums.len() >= 4 { nums[2].min(nums[3]) } else { 0.0 }
156}
157
158pub fn update_cluster_assignments(conn: &Connection, assignments: &[(i64, Option<i64>)]) -> rusqlite::Result<()> {
159    for (face_id, cluster_id) in assignments {
160        conn.execute(
161            "UPDATE faces SET cluster_id = ?1 WHERE id = ?2",
162            rusqlite::params![cluster_id, face_id],
163        )?;
164    }
165    Ok(())
166}
167
168pub fn hashes_with_faces(conn: &Connection) -> rusqlite::Result<Vec<String>> {
169    let mut stmt = conn.prepare("SELECT DISTINCT hash FROM faces ORDER BY hash")?;
170    let rows = stmt.query_map([], |r| r.get(0))?;
171    rows.collect()
172}
173
174/// (face_id, person_label, bbox) for one labeled face.
175pub type LabeledFace = (i64, String, String);
176
177/// Maps a file hash to every labeled face on it, as returned by
178/// `labeled_faces_by_hash`.
179pub type LabeledFacesByHash = HashMap<String, Vec<LabeledFace>>;
180
181/// Returns, for every hash that has at least one confirmed+labeled face, the
182/// list of (face_id, person_label, bbox) for that hash. One batched query
183/// covering every hash, not one query per file, safe to call once per
184/// report generation without N+1 overhead.
185pub fn labeled_faces_by_hash(conn: &Connection) -> rusqlite::Result<LabeledFacesByHash> {
186    let mut stmt = conn.prepare(
187        "SELECT hash, id, bbox, person_label FROM faces \
188         WHERE confirmed = 1 AND person_label IS NOT NULL \
189         ORDER BY hash, id",
190    )?;
191    let rows = stmt.query_map([], |r| {
192        Ok((
193            r.get::<_, String>(0)?,
194            r.get::<_, i64>(1)?,
195            r.get::<_, String>(2)?,
196            r.get::<_, String>(3)?,
197        ))
198    })?;
199    let mut map: LabeledFacesByHash = HashMap::new();
200    for row in rows {
201        let (hash, id, bbox, label) = row?;
202        map.entry(hash).or_default().push((id, label, bbox));
203    }
204    Ok(map)
205}
206
207#[cfg(test)]
208fn make_embedding(vals: &[f32]) -> Vec<u8> {
209    vals.iter().flat_map(|&v| f16::from_f32(v).to_le_bytes()).collect()
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    fn open() -> Connection {
217        let conn = Connection::open_in_memory().unwrap();
218        create_faces_table(&conn).unwrap();
219        conn
220    }
221
222    #[test]
223    fn create_table_idempotent() {
224        let conn = open();
225        create_faces_table(&conn).unwrap();
226    }
227
228    #[test]
229    fn insert_and_load_embedding() {
230        let conn = open();
231        let emb = make_embedding(&vec![0.5f32; 512]);
232        replace_faces_for_hash(&conn, "habc", &[FaceRow {
233            hash: "habc".into(), bbox: "0,0,50,50".into(), landmark: None,
234            embedding: emb, cluster_id: None, person_label: None, confirmed: 0, is_primary: 0,
235        }]).unwrap();
236        let rows = load_face_embeddings(&conn).unwrap();
237        assert_eq!(rows.len(), 1);
238        let (id, emb_f32) = &rows[0];
239        assert!(*id > 0);
240        assert_eq!(emb_f32.len(), 512);
241        assert!((emb_f32[0] - 0.5).abs() < 0.01);
242    }
243
244    #[test]
245    fn replace_removes_old_rows_for_same_hash() {
246        let conn = open();
247        let emb = make_embedding(&vec![0.0f32; 512]);
248        replace_faces_for_hash(&conn, "h1", &[
249            FaceRow { hash: "h1".into(), bbox: "0,0,10,10".into(), landmark: None, embedding: emb.clone(), cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 },
250            FaceRow { hash: "h1".into(), bbox: "20,0,10,10".into(), landmark: None, embedding: emb.clone(), cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 },
251        ]).unwrap();
252        replace_faces_for_hash(&conn, "h1", &[
253            FaceRow { hash: "h1".into(), bbox: "99,0,10,10".into(), landmark: None, embedding: emb, cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 },
254        ]).unwrap();
255        let rows = load_face_embeddings(&conn).unwrap();
256        assert_eq!(rows.len(), 1);
257    }
258
259    #[test]
260    fn update_cluster_assignments_works() {
261        let conn = open();
262        let emb = make_embedding(&vec![0.0f32; 512]);
263        replace_faces_for_hash(&conn, "h1", &[FaceRow { hash: "h1".into(), bbox: "0,0,10,10".into(), landmark: None, embedding: emb, cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 }]).unwrap();
264        let rows = load_face_embeddings(&conn).unwrap();
265        let id = rows[0].0;
266        update_cluster_assignments(&conn, &[(id, Some(3))]).unwrap();
267        let n: i64 = conn.query_row("SELECT cluster_id FROM faces WHERE id=?1", [id], |r| r.get(0)).unwrap();
268        assert_eq!(n, 3);
269    }
270
271    #[test]
272    fn load_faces_for_clustering_returns_bbox_min_side() {
273        let conn = open();
274        let emb = make_embedding(&vec![0.25f32; 512]);
275        // bbox "x,y,w,h": min side is min(w,h).
276        replace_faces_for_hash(&conn, "h1", &[
277            FaceRow { hash: "h1".into(), bbox: "10,10,200,300".into(), landmark: None, embedding: emb.clone(), cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 },
278            FaceRow { hash: "h1".into(), bbox: "0,0,40,25".into(), landmark: None, embedding: emb, cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 },
279        ]).unwrap();
280        let mut rows = load_faces_for_clustering(&conn).unwrap();
281        rows.sort_by(|a, b| b.2.total_cmp(&a.2));
282        assert_eq!(rows[0].2, 200.0, "min side of 200x300 bbox");
283        assert_eq!(rows[1].2, 25.0, "min side of 40x25 bbox");
284        assert_eq!(rows[0].1.len(), 512, "embedding still decoded");
285    }
286
287    #[test]
288    fn mark_scanned_records_hash_even_with_zero_faces() {
289        let conn = open();
290        // A hash processed with no detected faces leaves no `faces` row, but
291        // must still be recorded as scanned so it is not re-processed.
292        mark_scanned(&conn, "noface").unwrap();
293        assert_eq!(scanned_hashes(&conn).unwrap(), vec!["noface".to_string()]);
294        // hashes_with_faces stays empty, the marker is independent of faces.
295        assert!(hashes_with_faces(&conn).unwrap().is_empty());
296    }
297
298    #[test]
299    fn mark_scanned_is_idempotent() {
300        let conn = open();
301        mark_scanned(&conn, "h").unwrap();
302        mark_scanned(&conn, "h").unwrap();
303        assert_eq!(scanned_hashes(&conn).unwrap().len(), 1);
304    }
305
306    #[test]
307    fn select_unscanned_skips_dedups_and_limits() {
308        // Two paths share hash "a"; "b" is skipped; "c","d","e" remain.
309        let all = vec![
310            ("/1.jpg".to_string(), "a".to_string()),
311            ("/1copy.jpg".to_string(), "a".to_string()),
312            ("/2.jpg".to_string(), "b".to_string()),
313            ("/3.jpg".to_string(), "c".to_string()),
314            ("/4.jpg".to_string(), "d".to_string()),
315            ("/5.jpg".to_string(), "e".to_string()),
316        ];
317        let skip: std::collections::HashSet<String> = ["b".to_string()].into_iter().collect();
318        // No limit: one path per unscanned hash (a,c,d,e), b excluded.
319        let out = select_unscanned(&all, &skip, None);
320        assert_eq!(
321            out.iter().map(|(_, h)| h.clone()).collect::<Vec<_>>(),
322            vec!["a", "c", "d", "e"]
323        );
324        // Limit 2: first two unscanned hashes only.
325        let out2 = select_unscanned(&all, &skip, Some(2));
326        assert_eq!(
327            out2.iter().map(|(_, h)| h.clone()).collect::<Vec<_>>(),
328            vec!["a", "c"]
329        );
330    }
331
332    #[test]
333    fn hashes_with_faces_returns_inserted_hash() {
334        let conn = open();
335        let emb = make_embedding(&vec![0.0f32; 512]);
336        replace_faces_for_hash(&conn, "myhash", &[FaceRow { hash: "myhash".into(), bbox: "0,0,10,10".into(), landmark: None, embedding: emb, cluster_id: None, person_label: None, confirmed: 0, is_primary: 0 }]).unwrap();
337        let hashes = hashes_with_faces(&conn).unwrap();
338        assert_eq!(hashes, vec!["myhash"]);
339    }
340
341    #[test]
342    fn labeled_faces_by_hash_returns_only_confirmed_labeled() {
343        let conn = Connection::open_in_memory().unwrap();
344        create_faces_table(&conn).unwrap();
345        conn.execute_batch(
346            "INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
347             VALUES ('h1', '0,0,10,10', X'0000', 'Alice', 1); \
348             INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
349             VALUES ('h1', '20,20,10,10', X'0000', NULL, 0); \
350             INSERT INTO faces (hash, bbox, embedding, person_label, confirmed) \
351             VALUES ('h2', '0,0,10,10', X'0000', 'Bob', 1);",
352        )
353        .unwrap();
354
355        let map = labeled_faces_by_hash(&conn).unwrap();
356        assert_eq!(map.len(), 2, "expected two hashes with labeled faces");
357        let h1 = &map["h1"];
358        assert_eq!(h1.len(), 1, "unconfirmed/unlabeled face must be excluded");
359        assert_eq!(h1[0].1, "Alice");
360        assert_eq!(map["h2"][0].1, "Bob");
361    }
362}
videre_core/face_db.rs

videre_core/
face_db.rs
