nuviz_cli/data/
images.rs

1use std::path::{Path, PathBuf};
2
3/// An image file discovered in an experiment's images/ directory.
4#[derive(Debug, Clone)]
5pub struct ImageEntry {
6    pub path: PathBuf,
7    pub step: u64,
8    pub tag: String,
9    pub size_bytes: u64,
10}
11
12/// Discover all images in an experiment directory, sorted by step then tag.
13///
14/// Expected filename pattern: `step_NNNNNN_<tag>.png`
15pub fn discover_images(experiment_dir: &Path) -> Vec<ImageEntry> {
16    let images_dir = experiment_dir.join("images");
17    if !images_dir.is_dir() {
18        return Vec::new();
19    }
20
21    let mut entries = Vec::new();
22
23    let read_dir = match std::fs::read_dir(&images_dir) {
24        Ok(rd) => rd,
25        Err(_) => return Vec::new(),
26    };
27
28    for entry in read_dir.flatten() {
29        let path = entry.path();
30        if path.extension().and_then(|e| e.to_str()) != Some("png") {
31            continue;
32        }
33
34        if let Some(parsed) = parse_image_filename(&path) {
35            let size_bytes = entry.metadata().map(|m| m.len()).unwrap_or(0);
36            entries.push(ImageEntry {
37                path,
38                step: parsed.0,
39                tag: parsed.1,
40                size_bytes,
41            });
42        }
43    }
44
45    entries.sort_by(|a, b| a.step.cmp(&b.step).then_with(|| a.tag.cmp(&b.tag)));
46    entries
47}
48
49/// Find the latest image matching an optional tag filter.
50pub fn find_latest_image(experiment_dir: &Path, tag: Option<&str>) -> Option<ImageEntry> {
51    let mut images = discover_images(experiment_dir);
52    if let Some(tag_filter) = tag {
53        images.retain(|e| e.tag == tag_filter);
54    }
55    images.into_iter().last()
56}
57
58/// Parse `step_NNNNNN_<tag>.png` -> (step, tag)
59fn parse_image_filename(path: &Path) -> Option<(u64, String)> {
60    let stem = path.file_stem()?.to_str()?;
61    let parts: Vec<&str> = stem.splitn(3, '_').collect();
62    if parts.len() < 3 || parts[0] != "step" {
63        return None;
64    }
65    let step = parts[1].parse::<u64>().ok()?;
66    let tag = parts[2].to_string();
67    Some((step, tag))
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use std::fs;
74
75    #[test]
76    fn test_parse_image_filename() {
77        let path = Path::new("/tmp/images/step_000500_render.png");
78        let (step, tag) = parse_image_filename(path).unwrap();
79        assert_eq!(step, 500);
80        assert_eq!(tag, "render");
81    }
82
83    #[test]
84    fn test_parse_image_filename_with_underscores_in_tag() {
85        let path = Path::new("/tmp/images/step_001000_depth_map.png");
86        let (step, tag) = parse_image_filename(path).unwrap();
87        assert_eq!(step, 1000);
88        assert_eq!(tag, "depth_map");
89    }
90
91    #[test]
92    fn test_parse_image_filename_invalid() {
93        assert!(parse_image_filename(Path::new("/tmp/images/random.png")).is_none());
94        assert!(parse_image_filename(Path::new("/tmp/images/step_abc_tag.png")).is_none());
95    }
96
97    #[test]
98    fn test_discover_images() {
99        let dir = tempfile::tempdir().unwrap();
100        let images_dir = dir.path().join("images");
101        fs::create_dir_all(&images_dir).unwrap();
102
103        // Create test image files (just empty files for discovery)
104        fs::write(images_dir.join("step_000000_render.png"), b"fake").unwrap();
105        fs::write(images_dir.join("step_000000_gt.png"), b"fake").unwrap();
106        fs::write(images_dir.join("step_000100_render.png"), b"fake").unwrap();
107        fs::write(images_dir.join("not_an_image.txt"), b"nope").unwrap();
108
109        let entries = discover_images(dir.path());
110        assert_eq!(entries.len(), 3);
111        assert_eq!(entries[0].step, 0);
112        assert_eq!(entries[0].tag, "gt");
113        assert_eq!(entries[1].step, 0);
114        assert_eq!(entries[1].tag, "render");
115        assert_eq!(entries[2].step, 100);
116    }
117
118    #[test]
119    fn test_discover_images_empty() {
120        let dir = tempfile::tempdir().unwrap();
121        let entries = discover_images(dir.path());
122        assert!(entries.is_empty());
123    }
124
125    #[test]
126    fn test_find_latest_image() {
127        let dir = tempfile::tempdir().unwrap();
128        let images_dir = dir.path().join("images");
129        fs::create_dir_all(&images_dir).unwrap();
130
131        fs::write(images_dir.join("step_000000_render.png"), b"fake").unwrap();
132        fs::write(images_dir.join("step_000100_render.png"), b"fake").unwrap();
133        fs::write(images_dir.join("step_000100_depth.png"), b"fake").unwrap();
134
135        let latest = find_latest_image(dir.path(), Some("render")).unwrap();
136        assert_eq!(latest.step, 100);
137        assert_eq!(latest.tag, "render");
138
139        let latest_any = find_latest_image(dir.path(), None).unwrap();
140        assert_eq!(latest_any.step, 100);
141    }
142}
nuviz_cli/data/images.rs

nuviz_cli/data/
images.rs
