capture_utils/
lib.rs

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

capture_utils/
lib.rs
