flydra2/
mini_arenas.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{collections::BTreeMap, path::Path};
5
6use serde::{Deserialize, Serialize};
7
8use braid_types::MiniArenaConfig;
9
10use crate::{MyFloat, Result, bundled_data::BundledAllCamsOneFrameUndistorted};
11
12/// Into a mini arena.
13///
14/// This is different from [braid_types::MiniArenaLocator] because the mini
15/// arena must exist.
16///
17/// Newtype wrapper around [u8].
18#[derive(Debug, Clone, Copy, PartialEq)]
19pub(crate) struct MiniArenaIndex(u8);
20impl MiniArenaIndex {
21    pub(crate) fn new(val: u8) -> Self {
22        MiniArenaIndex(val)
23    }
24    pub(crate) fn idx(&self) -> usize {
25        self.0 as usize
26    }
27}
28
29const NO_MINI_ARENA_MARKER: u8 = 255;
30
31pub(crate) enum MiniArenaLocator {
32    /// Location is not possible.
33    OutOfBounds,
34    /// Not in mini arena.
35    NotInMiniArena,
36    /// In mini arena with index.
37    Index(MiniArenaIndex),
38    /// No mini arenas are in use.
39    OneArena,
40}
41
42/// Image of a mini arena for a calibrated camera.
43pub(crate) struct MiniArenaImage {
44    width: usize,
45    // height is data.len() / width
46    data: Vec<u8>,
47}
48
49impl std::fmt::Debug for MiniArenaImage {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("MiniArenaImage")
52            .field("width", &self.width)
53            .field("data.len()", &self.data.len())
54            .finish_non_exhaustive()
55    }
56}
57
58impl MiniArenaImage {
59    /// Get mini arena locator.
60    pub(crate) fn get_mini_arena(&self, x: usize, y: usize) -> MiniArenaLocator {
61        let idx = y * self.width + x;
62        match self.data.get(idx) {
63            None => MiniArenaLocator::OutOfBounds,
64            Some(&NO_MINI_ARENA_MARKER) => MiniArenaLocator::NotInMiniArena,
65            Some(idx) => MiniArenaLocator::Index(MiniArenaIndex::new(*idx)),
66        }
67    }
68}
69
70pub struct MiniArenaDebugConfig {
71    /// Directory to save mini arena images.
72    pub output_png_path: camino::Utf8PathBuf,
73    /// Background image to use for mini arena images.
74    pub background_image_jpeg_buf: Option<Vec<u8>>,
75    pub april_detections: Option<Vec<braid_apriltag_types::AprilTagCoords2D>>,
76}
77
78impl std::fmt::Debug for MiniArenaDebugConfig {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        f.debug_struct("MiniArenaDebugConfig")
81            .finish_non_exhaustive()
82    }
83}
84
85/// Build per-camera mini-arena images.
86pub(crate) fn build_mini_arena_images(
87    recon: Option<&flydra_mvg::FlydraMultiCameraSystem<MyFloat>>,
88    mini_arena_config: &MiniArenaConfig,
89    mini_arena_debug_cfg: Option<&MiniArenaDebugConfig>,
90) -> Result<BTreeMap<String, MiniArenaImage>> {
91    let mut mini_arena_images = BTreeMap::new();
92    let recon = match recon {
93        None => {
94            return Ok(mini_arena_images);
95        }
96        Some(recon) => recon,
97    };
98    match mini_arena_config {
99        MiniArenaConfig::NoMiniArena => {}
100        MiniArenaConfig::XYGrid(xy_grid_cfg) => {
101            for cam in recon.cameras() {
102                let sz = cam.width() * cam.height();
103                let mut mini_arena_image = vec![NO_MINI_ARENA_MARKER; sz];
104
105                for row in 0..cam.height() {
106                    for col in 0..cam.width() {
107                        let pt = braid_mvg::DistortedPixel {
108                            coords: nalgebra::Point2::new(col as f64, row as f64),
109                        };
110                        let ray = cam.project_distorted_pixel_to_ray(&pt);
111                        let coords_3d = crate::flat_2d::ray_to_flat_3d(&ray);
112                        if let Some(coords_3d) = coords_3d {
113                            let coords = [coords_3d.x, coords_3d.y, coords_3d.z];
114                            if let Some(arena_idx) = xy_grid_cfg.get_arena_index(&coords).idx() {
115                                let coords_idx = row * cam.width() + col;
116                                mini_arena_image[coords_idx] = arena_idx;
117                            }
118                        }
119                    }
120                }
121
122                if let Some(cfg) = &mini_arena_debug_cfg {
123                    // save debug image of mini arenas.
124                    use machine_vision_formats::pixel_format::Mono8;
125                    let well_image = machine_vision_formats::image_ref::ImageRef::<Mono8>::new(
126                        cam.width().try_into().unwrap(),
127                        cam.height().try_into().unwrap(),
128                        cam.width(),
129                        &mini_arena_image,
130                    )
131                    .unwrap();
132                    let well_jpeg_buf = convert_image::frame_to_encoded_buffer(
133                        &well_image,
134                        convert_image::EncoderOptions::Jpeg(100),
135                    )
136                    .unwrap();
137
138                    annotate_mini_arena_image(&well_jpeg_buf, cfg, &cam, xy_grid_cfg).unwrap();
139                }
140
141                mini_arena_images.insert(
142                    cam.name().to_string(),
143                    MiniArenaImage {
144                        width: cam.width(),
145                        data: mini_arena_image,
146                    },
147                );
148            }
149        }
150    }
151    Ok(mini_arena_images)
152}
153
154/// On top of background image, draw mini arena numbers.
155fn annotate_mini_arena_image(
156    well_jpeg_buf: &[u8],
157    cfg: &MiniArenaDebugConfig,
158    cam: &flydra_mvg::MultiCamera<f64>,
159    xy_grid_cfg: &braid_types::XYGridConfig,
160) -> eyre::Result<()> {
161    let svg_width = cam.width();
162    let svg_height = cam.height();
163    // Draw SVG
164    let mut wtr = tagger::new(tagger::upgrade_write(Vec::<u8>::new()));
165    // let svg_width = self.cum_width + n_pics * 2 * composite_margin_pixels;
166    // let svg_height = self.cum_height + 2 * composite_margin_pixels;
167    wtr.elem("svg", |d| {
168        d.attr("xmlns", "http://www.w3.org/2000/svg")?;
169        d.attr("xmlns:xlink", "http://www.w3.org/1999/xlink")?;
170        d.attr("viewBox", format_args!("0 0 {} {}", svg_width, svg_height))
171    })?
172    .build(|w| {
173        use base64::Engine;
174
175        // Draw background image
176        if let Some(jpeg_buf) = &cfg.background_image_jpeg_buf {
177            let jpeg_base64_buf = base64::engine::general_purpose::STANDARD.encode(jpeg_buf);
178            let data_url = format!("data:image/jpeg;base64,{}", jpeg_base64_buf);
179            w.single("image", |d| {
180                d.attr("x", 0)?;
181                d.attr("y", 0)?;
182                d.attr("width", svg_width)?;
183                d.attr("height", svg_height)?;
184                d.attr("xlink:href", data_url)
185            })?;
186        }
187
188        // Draw well image
189        {
190            let jpeg_base64_buf = base64::engine::general_purpose::STANDARD.encode(well_jpeg_buf);
191            let data_url = format!("data:image/jpeg;base64,{}", jpeg_base64_buf);
192            w.single("image", |d| {
193                d.attr("x", 0)?;
194                d.attr("y", 0)?;
195                d.attr("width", svg_width)?;
196                d.attr("height", svg_height)?;
197                d.attr("opacity", "0.3")?;
198                d.attr("xlink:href", data_url)
199            })?;
200        }
201
202        // Draw april tag detections
203        if let Some(detections) = &cfg.april_detections {
204            for detection in detections.iter() {
205
206                let id_str = if detection.vertical_flip {
207                    format!("{} (v)", detection.id)
208                } else {
209                    format!("{}", detection.id)
210                };
211                w.elem("text", |d| {
212                d.attr("x", format!("{}", detection.x))?;
213                d.attr("y", format!("{}", detection.y))?;
214                d.attr("text-anchor ", "middle")?;
215                d.attr("dominant-baseline", "middle")?;
216                d.attr(
217                    "style",
218                    "font-family: Arial, Helvetica, sans-serif; font-size: 30px; fill: red;",
219                )?;
220                Ok(())
221            })?
222            .build(|w| w.put_raw(id_str))?;
223            }
224        }
225
226        // Draw mini arena numbers.
227        for center in xy_grid_cfg.iter_centers() {
228            let coord = [center.0, center.1, 0.0];
229            // Get mini arena index.
230            let idx = xy_grid_cfg.get_arena_index(&coord).idx().unwrap();
231
232            // The arena center is in 3D world coordinates. Project it to pixel
233            // coordinates.
234            let pt = cam
235                .project_3d_to_distorted_pixel(&braid_mvg::PointWorldFrame::from(coord))
236                .coords;
237
238            w.elem("text", |d| {
239                d.attr("x", format!("{}", pt.x))?;
240                d.attr("y", format!("{}", pt.y))?;
241                d.attr("text-anchor ", "middle")?;
242                d.attr("dominant-baseline", "middle")?;
243                d.attr(
244                    "style",
245                    "font-family: Arial, Helvetica, sans-serif; font-size: 40px; fill: deepskyblue;",
246                )?;
247                Ok(())
248            })?
249            .build(|w| w.put_raw(format!("{}", idx)))?;
250        }
251
252        Ok(())
253    })?;
254    // Get the SVG file contents.
255    let fmt_wtr = wtr.into_writer();
256    let svg_buf = {
257        fmt_wtr.error?;
258        fmt_wtr.inner
259    };
260
261    let mut usvg_opt = usvg::Options::default();
262    usvg_opt.fontdb_mut().load_system_fonts();
263
264    // Now parse the SVG file.
265    let rtree = usvg::Tree::from_data(&svg_buf, &usvg_opt)?;
266    // Now render the SVG file to a pixmap.
267    let pixmap_size = rtree.size().to_int_size();
268    let mut pixmap =
269        resvg::tiny_skia::Pixmap::new(pixmap_size.width(), pixmap_size.height()).unwrap();
270    resvg::render(
271        &rtree,
272        resvg::tiny_skia::Transform::default(),
273        &mut pixmap.as_mut(),
274    );
275
276    pixmap.save_png(&cfg.output_png_path)?;
277    tracing::info!("Saved well image to {}", cfg.output_png_path);
278
279    Ok(())
280}
281
282// ------ debug to CSV stuff ---------------
283
284/// Debugging structure to save to CSV files
285#[derive(Clone, Debug, Serialize, Deserialize)]
286struct MiniArenaPointPerCamFlat {
287    frame: u64,
288    cam_name: String,
289    mini_arena_idx: usize,
290    undistorted_x: f64,
291    undistorted_y: f64,
292    distorted_idx: u8,
293    distorted_x: f64,
294    distorted_y: f64,
295}
296
297impl MiniArenaPointPerCamFlat {
298    fn new(
299        frame: u64,
300        cam_name: String,
301        mini_arena_idx: usize,
302        orig: &crate::bundled_data::MiniArenaPointPerCam,
303    ) -> Self {
304        Self {
305            frame,
306            cam_name,
307            mini_arena_idx,
308            undistorted_x: orig.undistorted.x,
309            undistorted_y: orig.undistorted.y,
310            distorted_idx: orig.numbered_raw_udp_point.idx,
311            distorted_x: orig.numbered_raw_udp_point.pt.x0_abs,
312            distorted_y: orig.numbered_raw_udp_point.pt.y0_abs,
313        }
314    }
315}
316
317// TODO: this gets called from an async task but does blocking IO. It should be
318// rewritten to use async IO.
319pub(crate) struct MiniArenaAssignmentDebug {
320    wtr: csv::Writer<std::io::BufWriter<std::fs::File>>,
321}
322
323impl MiniArenaAssignmentDebug {
324    pub(crate) fn new<P: AsRef<Path>>(path: P) -> crate::Result<Self> {
325        let fd = std::fs::File::create(path)?;
326        let bufwriter = std::io::BufWriter::new(fd);
327        let wtr = csv::Writer::from_writer(bufwriter);
328        Ok(Self { wtr })
329    }
330
331    pub(crate) fn write_frame(
332        &mut self,
333        undistorted: &BundledAllCamsOneFrameUndistorted,
334    ) -> crate::Result<()> {
335        let frame = undistorted.tdpt.frame.0;
336
337        for (mini_arena_idx, mini_arena) in undistorted.per_mini_arena.iter().enumerate() {
338            for (mini_arena_cam, mini_arena_data) in mini_arena.per_cam.iter() {
339                for orig in mini_arena_data.iter() {
340                    let row = MiniArenaPointPerCamFlat::new(
341                        frame,
342                        mini_arena_cam.as_str().to_string(),
343                        mini_arena_idx,
344                        orig,
345                    );
346
347                    self.wtr.serialize(row)?;
348                }
349            }
350        }
351
352        Ok(())
353    }
354}
flydra2/mini_arenas.rs

flydra2/
mini_arenas.rs
