capture_utils/
lib.rs

1//! Filesystem-based recorder implementations and dataset management utilities.
2//!
3//! This crate provides:
4//! - `JsonRecorder`: Default recorder that writes frame metadata to JSON files.
5//! - Overlay generation from label JSON files (draws bounding boxes on images).
6//! - Dataset pruning utilities for filtering and copying runs.
7//!
8//! Integrates with `vision_core::interfaces::Recorder` trait and `data_contracts` schemas.
9
10use data_contracts::capture::CaptureMetadata;
11use image::Rgba;
12use std::fs;
13use std::io::{BufWriter, Write};
14use std::path::{Path, PathBuf};
15use std::time::{SystemTime, UNIX_EPOCH};
16use vision_core::prelude::{FrameRecord, Label, Recorder};
17
18/// Default file-based recorder: writes frame metadata/labels to `run_dir/labels/frame_XXXXX.json`.
19pub struct JsonRecorder {
20    pub run_dir: PathBuf,
21}
22
23impl JsonRecorder {
24    pub fn new(run_dir: impl Into<PathBuf>) -> Self {
25        Self {
26            run_dir: run_dir.into(),
27        }
28    }
29}
30
31impl Recorder for JsonRecorder {
32    fn record(&mut self, record: &FrameRecord) -> std::io::Result<()> {
33        let labels_dir = self.run_dir.join("labels");
34        fs::create_dir_all(&labels_dir)?;
35        let image = record
36            .frame
37            .path
38            .as_ref()
39            .map(|p| p.to_string_lossy().into_owned())
40            .unwrap_or_else(|| format!("frame_{:05}.png", record.frame.id));
41        let unix_time = SystemTime::now()
42            .duration_since(UNIX_EPOCH)
43            .map(|d| d.as_secs_f64())
44            .unwrap_or(0.0);
45        let meta = build_capture_metadata(record, unix_time, image);
46        meta.validate()
47            .map_err(|e| std::io::Error::other(format!("validation failed: {e}")))?;
48        let out = labels_dir.join(format!("frame_{:05}.json", record.frame.id));
49        let mut writer = BufWriter::new(fs::File::create(out)?);
50        serde_json::to_writer_pretty(&mut writer, &meta)?;
51        writer.write_all(b"\n")?;
52        Ok(())
53    }
54}
55
56pub fn build_capture_metadata(
57    record: &FrameRecord,
58    unix_time: f64,
59    image: String,
60) -> CaptureMetadata {
61    CaptureMetadata {
62        frame_id: record.frame.id,
63        sim_time: record.frame.timestamp,
64        unix_time,
65        image,
66        image_present: record.frame.path.is_some(),
67        camera_active: record.camera_active,
68        label_seed: record.label_seed,
69        labels: record.labels.to_vec(),
70    }
71}
72
73/// Helper for inference outputs: callers set label provenance on `Label` before recording.
74pub fn build_inference_metadata(
75    frame: vision_core::prelude::Frame,
76    labels: &[Label],
77    camera_active: bool,
78    label_seed: u64,
79    unix_time: f64,
80) -> CaptureMetadata {
81    let image = frame
82        .path
83        .as_ref()
84        .map(|p| p.to_string_lossy().into_owned())
85        .unwrap_or_else(|| format!("frame_{:05}.png", frame.id));
86    let record = FrameRecord {
87        frame,
88        labels,
89        camera_active,
90        label_seed,
91    };
92    build_capture_metadata(&record, unix_time, image)
93}
94
95/// Generate overlay PNGs from label JSONs in a run directory.
96pub fn generate_overlays(run_dir: &Path) -> anyhow::Result<()> {
97    let labels_dir = run_dir.join("labels");
98    let out_dir = run_dir.join("overlays");
99    fs::create_dir_all(&out_dir)?;
100
101    for entry in fs::read_dir(&labels_dir).into_iter().flatten() {
102        let Ok(entry) = entry else { continue };
103        let path = entry.path();
104        if path.extension().and_then(|e| e.to_str()) != Some("json") {
105            continue;
106        }
107        let Ok(bytes) = fs::read(&path) else { continue };
108        let Ok(meta) = serde_json::from_slice::<CaptureMetadata>(&bytes) else {
109            continue;
110        };
111        if !meta.image_present {
112            continue;
113        }
114        let img_path = run_dir.join(&meta.image);
115        if !img_path.exists() {
116            continue;
117        }
118        let Ok(mut img) = image::open(&img_path).map(|im| im.into_rgba8()) else {
119            continue;
120        };
121        let (w, h) = img.dimensions();
122        let clamp =
123            |v: f32, max: u32| -> u32 { v.max(0.0).min((max.saturating_sub(1)) as f32) as u32 };
124        for label in meta.labels.iter().filter_map(|l| l.bbox_px) {
125            let bbox_px = [
126                clamp(label[0], w),
127                clamp(label[1], h),
128                clamp(label[2], w),
129                clamp(label[3], h),
130            ];
131            draw_rect(&mut img, bbox_px, Rgba([255, 64, 192, 255]), 2);
132        }
133        let filename = Path::new(&meta.image)
134            .file_name()
135            .map(|s| s.to_string_lossy().into_owned())
136            .unwrap_or(meta.image);
137        let _ = img.save(out_dir.join(filename));
138    }
139    Ok(())
140}
141
142/// Prune a run directory into a destination root, copying kept artifacts.
143pub fn prune_run(input_run: &Path, output_root: &Path) -> std::io::Result<(usize, usize)> {
144    let run_name = input_run
145        .file_name()
146        .ok_or_else(|| std::io::Error::other("invalid run dir"))?;
147    let out_run = output_root.join(run_name);
148    fs::create_dir_all(out_run.join("labels"))?;
149    fs::create_dir_all(out_run.join("images"))?;
150    fs::create_dir_all(out_run.join("overlays"))?;
151
152    let manifest_in = input_run.join("run_manifest.json");
153    if manifest_in.exists() {
154        let _ = fs::copy(&manifest_in, out_run.join("run_manifest.json"));
155    }
156
157    for entry in fs::read_dir(input_run)? {
158        let entry = entry?;
159        let path = entry.path();
160        let name = entry.file_name();
161        let Some(fname) = name.to_str() else { continue };
162        if fname == "run_manifest.json" {
163            continue;
164        }
165        if path.is_dir() {
166            if let Some(dir_name) = path.file_name() {
167                let out_dir = out_run.join(dir_name);
168                fs::create_dir_all(&out_dir)?;
169                for f in fs::read_dir(&path)? {
170                    let f = f?;
171                    let out_f = out_dir.join(f.file_name());
172                    let _ = fs::copy(f.path(), out_f);
173                }
174            }
175            continue;
176        }
177    }
178
179    let kept = fs::read_dir(out_run.join("labels"))
180        .into_iter()
181        .flatten()
182        .count();
183    let skipped = fs::read_dir(input_run.join("labels"))
184        .into_iter()
185        .flatten()
186        .count()
187        .saturating_sub(kept);
188    Ok((kept, skipped))
189}
190
191fn draw_rect(img: &mut image::RgbaImage, bbox: [u32; 4], color: Rgba<u8>, thickness: u32) {
192    let (x0, y0, x1, y1) = (bbox[0], bbox[1], bbox[2], bbox[3]);
193    for t in 0..thickness {
194        for x in x0.saturating_sub(t)..=x1.saturating_add(t) {
195            if y0 >= t {
196                if let Some(p) = img.get_pixel_mut_checked(x, y0 - t) {
197                    *p = color;
198                }
199            }
200            if let Some(p) = img.get_pixel_mut_checked(x, y1 + t) {
201                *p = color;
202            }
203        }
204        for y in y0.saturating_sub(t)..=y1.saturating_add(t) {
205            if x0 >= t {
206                if let Some(p) = img.get_pixel_mut_checked(x0 - t, y) {
207                    *p = color;
208                }
209            }
210            if let Some(p) = img.get_pixel_mut_checked(x1 + t, y) {
211                *p = color;
212            }
213        }
214    }
215}
216
217#[allow(dead_code)]
218trait GetPixelChecked {
219    fn get_pixel_mut_checked(&mut self, x: u32, y: u32) -> Option<&mut Rgba<u8>>;
220}
221
222impl GetPixelChecked for image::RgbaImage {
223    fn get_pixel_mut_checked(&mut self, x: u32, y: u32) -> Option<&mut Rgba<u8>> {
224        if x < self.width() && y < self.height() {
225            Some(self.get_pixel_mut(x, y))
226        } else {
227            None
228        }
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use vision_core::prelude::{Frame, FrameRecord};
236
237    #[test]
238    fn json_recorder_writes_label() {
239        let dir = tempfile::tempdir().unwrap();
240        let run_dir = dir.path();
241        let mut recorder = JsonRecorder::new(run_dir);
242        let frame = Frame {
243            id: 1,
244            timestamp: 0.5,
245            rgba: None,
246            size: (640, 480),
247            path: Some(PathBuf::from("images/frame_00001.png")),
248        };
249        let record = FrameRecord {
250            frame,
251            labels: &[],
252            camera_active: true,
253            label_seed: 42,
254        };
255        recorder.record(&record).expect("write label");
256        let label_path = run_dir.join("labels/frame_00001.json");
257        assert!(label_path.exists(), "label file should be written");
258    }
259
260    #[test]
261    fn prune_run_copies_manifest_and_dirs() {
262        let dir = tempfile::tempdir().unwrap();
263        let input = dir.path().join("run_123");
264        let labels = input.join("labels");
265        let images = input.join("images");
266        let overlays = input.join("overlays");
267        fs::create_dir_all(&labels).unwrap();
268        fs::create_dir_all(&images).unwrap();
269        fs::create_dir_all(&overlays).unwrap();
270        fs::write(input.join("run_manifest.json"), "{}").unwrap();
271        fs::write(labels.join("frame_00001.json"), "{}").unwrap();
272        fs::write(images.join("frame_00001.png"), []).unwrap();
273        fs::write(overlays.join("frame_00001.png"), []).unwrap();
274
275        let out_root = dir.path().join("out");
276        let (kept, _skipped) = prune_run(&input, &out_root).expect("prune run");
277        assert_eq!(kept, 1);
278        assert!(out_root.join("run_123/run_manifest.json").exists());
279        assert!(out_root.join("run_123/labels/frame_00001.json").exists());
280        assert!(out_root.join("run_123/images/frame_00001.png").exists());
281        assert!(out_root.join("run_123/overlays/frame_00001.png").exists());
282    }
283}
capture_utils/lib.rs

capture_utils/
lib.rs
