Skip to main content

oxihuman_morph/
pose_database.rs

1//! Searchable pose database with metadata and similarity search.
2
3#[allow(dead_code)]
4pub struct PoseEntry {
5    pub id: u64,
6    pub name: String,
7    pub tags: Vec<String>,
8    pub joints: Vec<[f32; 4]>,
9    pub metadata: Vec<(String, String)>,
10    pub thumbnail_hash: u64,
11}
12
13#[allow(dead_code)]
14pub struct PoseDatabase {
15    pub entries: Vec<PoseEntry>,
16    pub next_id: u64,
17}
18
19#[allow(dead_code)]
20pub fn new_pose_database() -> PoseDatabase {
21    PoseDatabase {
22        entries: Vec::new(),
23        next_id: 1,
24    }
25}
26
27#[allow(dead_code)]
28pub fn add_pose_entry(
29    db: &mut PoseDatabase,
30    name: &str,
31    joints: Vec<[f32; 4]>,
32    tags: Vec<String>,
33) -> u64 {
34    let id = db.next_id;
35    db.next_id += 1;
36    db.entries.push(PoseEntry {
37        id,
38        name: name.to_string(),
39        tags,
40        joints,
41        metadata: Vec::new(),
42        thumbnail_hash: 0,
43    });
44    id
45}
46
47#[allow(dead_code)]
48pub fn get_pose(db: &PoseDatabase, id: u64) -> Option<&PoseEntry> {
49    db.entries.iter().find(|e| e.id == id)
50}
51
52#[allow(dead_code)]
53pub fn search_by_name<'a>(db: &'a PoseDatabase, query: &str) -> Vec<&'a PoseEntry> {
54    let query_lower = query.to_lowercase();
55    db.entries
56        .iter()
57        .filter(|e| e.name.to_lowercase().contains(&query_lower))
58        .collect()
59}
60
61#[allow(dead_code)]
62pub fn search_by_tag<'a>(db: &'a PoseDatabase, tag: &str) -> Vec<&'a PoseEntry> {
63    let tag_lower = tag.to_lowercase();
64    db.entries
65        .iter()
66        .filter(|e| e.tags.iter().any(|t| t.to_lowercase() == tag_lower))
67        .collect()
68}
69
70#[allow(dead_code)]
71pub fn pose_similarity(a: &PoseEntry, b: &PoseEntry) -> f32 {
72    let len = a.joints.len().min(b.joints.len());
73    if len == 0 {
74        return 0.0;
75    }
76    let mut dot = 0.0_f32;
77    let mut norm_a = 0.0_f32;
78    let mut norm_b = 0.0_f32;
79    for i in 0..len {
80        for k in 0..4 {
81            dot += a.joints[i][k] * b.joints[i][k];
82            norm_a += a.joints[i][k] * a.joints[i][k];
83            norm_b += b.joints[i][k] * b.joints[i][k];
84        }
85    }
86    let denom = norm_a.sqrt() * norm_b.sqrt();
87    if denom < 1e-9 {
88        0.0
89    } else {
90        (dot / denom).clamp(-1.0, 1.0)
91    }
92}
93
94#[allow(dead_code)]
95pub fn nearest_pose<'a>(db: &'a PoseDatabase, query_joints: &[[f32; 4]]) -> Option<&'a PoseEntry> {
96    if db.entries.is_empty() {
97        return None;
98    }
99    let query_entry = PoseEntry {
100        id: 0,
101        name: String::new(),
102        tags: Vec::new(),
103        joints: query_joints.to_vec(),
104        metadata: Vec::new(),
105        thumbnail_hash: 0,
106    };
107    db.entries.iter().max_by(|a, b| {
108        let sa = pose_similarity(a, &query_entry);
109        let sb = pose_similarity(b, &query_entry);
110        sa.partial_cmp(&sb).unwrap_or(std::cmp::Ordering::Equal)
111    })
112}
113
114#[allow(dead_code)]
115pub fn remove_pose(db: &mut PoseDatabase, id: u64) -> bool {
116    let before = db.entries.len();
117    db.entries.retain(|e| e.id != id);
118    db.entries.len() < before
119}
120
121#[allow(dead_code)]
122pub fn pose_count(db: &PoseDatabase) -> usize {
123    db.entries.len()
124}
125
126#[allow(dead_code)]
127pub fn all_tags(db: &PoseDatabase) -> Vec<&str> {
128    let mut seen = std::collections::HashSet::new();
129    let mut result = Vec::new();
130    for entry in &db.entries {
131        for tag in &entry.tags {
132            if seen.insert(tag.as_str()) {
133                result.push(tag.as_str());
134            }
135        }
136    }
137    result
138}
139
140#[allow(dead_code)]
141pub fn pose_database_to_json(db: &PoseDatabase) -> String {
142    let mut out = String::from("{\"entries\":[");
143    for (i, entry) in db.entries.iter().enumerate() {
144        if i > 0 {
145            out.push(',');
146        }
147        out.push_str(&format!(
148            "{{\"id\":{},\"name\":\"{}\",\"tags\":[{}]}}",
149            entry.id,
150            entry.name.replace('"', "\\\""),
151            entry
152                .tags
153                .iter()
154                .map(|t| format!("\"{}\"", t.replace('"', "\\\"")))
155                .collect::<Vec<_>>()
156                .join(",")
157        ));
158    }
159    out.push_str("]}");
160    out
161}
162
163#[allow(dead_code)]
164pub fn import_poses(db: &mut PoseDatabase, entries: Vec<PoseEntry>) {
165    for mut entry in entries {
166        entry.id = db.next_id;
167        db.next_id += 1;
168        db.entries.push(entry);
169    }
170}
171
172#[allow(dead_code)]
173pub fn sort_by_name(db: &mut PoseDatabase) {
174    db.entries.sort_by(|a, b| a.name.cmp(&b.name));
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    fn make_joints(v: f32) -> Vec<[f32; 4]> {
182        vec![[v, 0.0, 0.0, 0.0], [0.0, v, 0.0, 0.0]]
183    }
184
185    #[test]
186    fn test_new_pose_database() {
187        let db = new_pose_database();
188        assert_eq!(pose_count(&db), 0);
189        assert_eq!(db.next_id, 1);
190    }
191
192    #[test]
193    fn test_add_pose_entry() {
194        let mut db = new_pose_database();
195        let id = add_pose_entry(
196            &mut db,
197            "T-Pose",
198            make_joints(1.0),
199            vec!["idle".to_string()],
200        );
201        assert_eq!(id, 1);
202        assert_eq!(pose_count(&db), 1);
203    }
204
205    #[test]
206    fn test_get_pose() {
207        let mut db = new_pose_database();
208        let id = add_pose_entry(&mut db, "Walk", make_joints(1.0), vec![]);
209        let entry = get_pose(&db, id);
210        assert!(entry.is_some());
211        assert_eq!(entry.expect("should succeed").name, "Walk");
212    }
213
214    #[test]
215    fn test_get_pose_not_found() {
216        let db = new_pose_database();
217        assert!(get_pose(&db, 99).is_none());
218    }
219
220    #[test]
221    fn test_search_by_name() {
222        let mut db = new_pose_database();
223        add_pose_entry(&mut db, "T-Pose", make_joints(1.0), vec![]);
224        add_pose_entry(&mut db, "Walk Cycle", make_joints(0.5), vec![]);
225        let results = search_by_name(&db, "walk");
226        assert_eq!(results.len(), 1);
227        assert_eq!(results[0].name, "Walk Cycle");
228    }
229
230    #[test]
231    fn test_search_by_name_case_insensitive() {
232        let mut db = new_pose_database();
233        add_pose_entry(&mut db, "Running Pose", make_joints(1.0), vec![]);
234        let results = search_by_name(&db, "RUNNING");
235        assert_eq!(results.len(), 1);
236    }
237
238    #[test]
239    fn test_search_by_tag() {
240        let mut db = new_pose_database();
241        add_pose_entry(
242            &mut db,
243            "T-Pose",
244            make_joints(1.0),
245            vec!["idle".to_string()],
246        );
247        add_pose_entry(&mut db, "Run", make_joints(0.5), vec!["motion".to_string()]);
248        let results = search_by_tag(&db, "idle");
249        assert_eq!(results.len(), 1);
250        assert_eq!(results[0].name, "T-Pose");
251    }
252
253    #[test]
254    fn test_pose_similarity_identical() {
255        let joints = make_joints(1.0);
256        let entry_a = PoseEntry {
257            id: 1,
258            name: "A".to_string(),
259            tags: vec![],
260            joints: joints.clone(),
261            metadata: vec![],
262            thumbnail_hash: 0,
263        };
264        let entry_b = PoseEntry {
265            id: 2,
266            name: "B".to_string(),
267            tags: vec![],
268            joints,
269            metadata: vec![],
270            thumbnail_hash: 0,
271        };
272        let sim = pose_similarity(&entry_a, &entry_b);
273        assert!((sim - 1.0).abs() < 1e-5);
274    }
275
276    #[test]
277    fn test_pose_similarity_different() {
278        let entry_a = PoseEntry {
279            id: 1,
280            name: "A".to_string(),
281            tags: vec![],
282            joints: vec![[1.0, 0.0, 0.0, 0.0]],
283            metadata: vec![],
284            thumbnail_hash: 0,
285        };
286        let entry_b = PoseEntry {
287            id: 2,
288            name: "B".to_string(),
289            tags: vec![],
290            joints: vec![[-1.0, 0.0, 0.0, 0.0]],
291            metadata: vec![],
292            thumbnail_hash: 0,
293        };
294        let sim = pose_similarity(&entry_a, &entry_b);
295        assert!(sim < 0.0);
296    }
297
298    #[test]
299    fn test_nearest_pose() {
300        let mut db = new_pose_database();
301        add_pose_entry(&mut db, "A", vec![[1.0, 0.0, 0.0, 0.0]], vec![]);
302        add_pose_entry(&mut db, "B", vec![[0.0, 1.0, 0.0, 0.0]], vec![]);
303        let query = vec![[1.0, 0.0, 0.0, 0.0]];
304        let nearest = nearest_pose(&db, &query);
305        assert!(nearest.is_some());
306        assert_eq!(nearest.expect("should succeed").name, "A");
307    }
308
309    #[test]
310    fn test_nearest_pose_empty() {
311        let db = new_pose_database();
312        assert!(nearest_pose(&db, &[[1.0, 0.0, 0.0, 0.0]]).is_none());
313    }
314
315    #[test]
316    fn test_remove_pose() {
317        let mut db = new_pose_database();
318        let id = add_pose_entry(&mut db, "Test", make_joints(1.0), vec![]);
319        assert!(remove_pose(&mut db, id));
320        assert_eq!(pose_count(&db), 0);
321        assert!(!remove_pose(&mut db, id));
322    }
323
324    #[test]
325    fn test_pose_count() {
326        let mut db = new_pose_database();
327        assert_eq!(pose_count(&db), 0);
328        add_pose_entry(&mut db, "A", make_joints(1.0), vec![]);
329        add_pose_entry(&mut db, "B", make_joints(0.5), vec![]);
330        assert_eq!(pose_count(&db), 2);
331    }
332
333    #[test]
334    fn test_all_tags() {
335        let mut db = new_pose_database();
336        add_pose_entry(
337            &mut db,
338            "A",
339            make_joints(1.0),
340            vec!["idle".to_string(), "standing".to_string()],
341        );
342        add_pose_entry(
343            &mut db,
344            "B",
345            make_joints(0.5),
346            vec!["idle".to_string(), "motion".to_string()],
347        );
348        let tags = all_tags(&db);
349        assert_eq!(tags.len(), 3);
350        assert!(tags.contains(&"idle"));
351        assert!(tags.contains(&"standing"));
352        assert!(tags.contains(&"motion"));
353    }
354
355    #[test]
356    fn test_sort_by_name() {
357        let mut db = new_pose_database();
358        add_pose_entry(&mut db, "Zebra", make_joints(1.0), vec![]);
359        add_pose_entry(&mut db, "Alpha", make_joints(0.5), vec![]);
360        sort_by_name(&mut db);
361        assert_eq!(db.entries[0].name, "Alpha");
362        assert_eq!(db.entries[1].name, "Zebra");
363    }
364
365    #[test]
366    fn test_pose_database_to_json() {
367        let mut db = new_pose_database();
368        add_pose_entry(&mut db, "Test", make_joints(1.0), vec!["tag1".to_string()]);
369        let json = pose_database_to_json(&db);
370        assert!(json.contains("\"name\":\"Test\""));
371        assert!(json.contains("\"tag1\""));
372    }
373
374    #[test]
375    fn test_import_poses() {
376        let mut db = new_pose_database();
377        let entry = PoseEntry {
378            id: 999,
379            name: "Imported".to_string(),
380            tags: vec![],
381            joints: make_joints(1.0),
382            metadata: vec![],
383            thumbnail_hash: 0,
384        };
385        import_poses(&mut db, vec![entry]);
386        assert_eq!(pose_count(&db), 1);
387        assert_eq!(db.entries[0].id, 1);
388    }
389}