capture_utils/
lib.rs

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

capture_utils/
lib.rs
