adder_codec_rs/utils/
viz.rs

1use adder_codec_core::{Event, PixelAddress};
2#[cfg(feature = "open-cv")]
3use opencv::core::{Mat, MatTraitConst, MatTraitConstManual};
4use std::error::Error;
5use std::fs::File;
6use std::io;
7#[cfg(feature = "open-cv")]
8use std::io::BufWriter;
9use std::io::{Cursor, Write};
10use std::path::Path;
11use std::process::{Command, Output};
12use video_rs_adder_dep::{Frame};
13
14#[cfg(feature = "open-cv")]
15/// Writes a given [`Mat`] to a file
16/// # Errors
17/// * [`io::Error`] if there is an error writing to the file
18/// * [`opencv::Error`] if the [`Mat`] is malformed
19/// # Safety
20/// This function is unsafe because it calls `Mat::at_unchecked()` which is unsafe
21/// # Panics
22/// This function panics if the amount data written to the file is not equal to the amount of data
23/// in the [`Mat`].
24pub fn write_frame_to_video_cv(
25    frame: &Mat,
26    video_writer: &mut BufWriter<File>,
27) -> Result<(), Box<dyn Error>> {
28    let frame_size = frame.size()?;
29    let len = frame_size.width * frame_size.height * frame.channels();
30
31    // SAFETY:
32    // `frame` is a valid `Mat` and `len` is the number of elements in the `Mat`
33    unsafe {
34        for idx in 0..len {
35            let val: *const u8 = frame.at_unchecked(idx)? as *const u8;
36            let bytes_written = video_writer.write(std::slice::from_raw_parts(val, 1))?;
37            assert_eq!(bytes_written, 1);
38        }
39    }
40    Ok(())
41}
42
43/// Convenience function for converting binary grayscale data to an mp4. Used for testing.
44/// # Errors
45/// * [`io::Error`] if there is an error writing to the file
46pub fn encode_video_ffmpeg(raw_path: &str, video_path: &str) -> io::Result<Output> {
47    // ffmpeg -f rawvideo -pix_fmt gray -s:v 346x260 -r 60 -i ./tmp.gray8 -crf 0 -c:v libx264 ./output_file.mp4
48    println!("Writing reconstruction as .mp4 with ffmpeg");
49    Command::new("ffmpeg")
50        .args([
51            "-f", "rawvideo", "-pix_fmt", "gray", "-s:v", "346x260", "-r", "30", "-i", raw_path,
52            "-crf", "0", "-c:v", "libx264", "-y", video_path,
53        ])
54        .output()
55}
56
57#[allow(dead_code)]
58/// Convenience function for converting downloading a file at the given `video_url`. Used for testing.
59/// # Errors
60/// * [`io::Error`] if there is an error downloading or writing the file
61pub async fn download_file(
62    store_path: &str,
63    video_url: &str,
64) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
65    // Download the video example, if you don't already have it
66    let path_str = store_path;
67    if !Path::new(path_str).exists() {
68        let resp = reqwest::get(video_url).await?;
69        let mut file_out = File::create(path_str)?;
70        let mut data_in = Cursor::new(resp.bytes().await?);
71        std::io::copy(&mut data_in, &mut file_out)?;
72    }
73    Ok(())
74}
75
76/// The display mode for visualizing detected features
77#[derive(Debug, PartialEq, Clone, Copy)]
78pub enum ShowFeatureMode {
79    /// Don't show features at all
80    Off,
81
82    /// Show the feature only at the instant in which the pixel **becomes** a feature
83    Instant,
84
85    /// Show the feature until it's no longer a feature
86    Hold,
87}
88
89/// Assuming the given event is a feature, draw it on the given `img` as a white cross
90pub fn draw_feature_event(e: &Event, img: &mut Frame) {
91    draw_feature_coord(e.coord.x, e.coord.y, img, false, None)
92}
93
94/// Draw a cross on the given `img` at the given `x` and `y` coordinates
95pub fn draw_feature_coord(
96    x: PixelAddress,
97    y: PixelAddress,
98    img: &mut Frame,
99    three_color: bool,
100    color: Option<[u8; 3]>,
101) {
102    let draw_color: [u8; 3] = color.unwrap_or_else(|| [255, 255, 255]);
103
104    let radius = 2;
105
106    unsafe {
107        if three_color {
108            for i in -radius..=radius {
109                for c in 0..3 {
110                    *img.uget_mut(((y as i32 + i) as usize, (x as i32) as usize, c)) =
111                        draw_color[c];
112                    *img.uget_mut(((y as i32) as usize, (x as i32 + i) as usize, c)) =
113                        draw_color[c];
114                }
115            }
116        } else {
117            for i in -radius..=radius {
118                *img.uget_mut(((y as i32 + i) as usize, (x as i32) as usize, 0)) = draw_color[0];
119                *img.uget_mut(((y as i32) as usize, (x as i32 + i) as usize, 0)) = draw_color[0];
120            }
121        }
122    }
123}
124
125/// Draw a rectangle on the given `img` with the given `color`
126pub fn draw_rect(
127    x1: PixelAddress,
128    y1: PixelAddress,
129    x2: PixelAddress,
130    y2: PixelAddress,
131    img: &mut Frame,
132    three_color: bool,
133    color: Option<[u8; 3]>,
134) {
135    let draw_color: [u8; 3] = color.unwrap_or_else(|| [255, 255, 255]);
136
137    unsafe {
138        if three_color {
139            for i in x1..=x2 {
140                for c in 0..3 {
141                    *img.uget_mut((y1 as usize, i as usize, c)) = draw_color[c];
142                    *img.uget_mut((y2 as usize, i as usize, c)) = draw_color[c];
143                }
144            }
145            for i in y1..=y2 {
146                for c in 0..3 {
147                    *img.uget_mut((i as usize, x1 as usize, c)) = draw_color[c];
148                    *img.uget_mut((i as usize, x2 as usize, c)) = draw_color[c];
149                }
150            }
151        } else {
152            for i in x1..=x2 {
153                *img.uget_mut((y1 as usize, i as usize, 0)) = draw_color[0];
154                *img.uget_mut((y2 as usize, i as usize, 0)) = draw_color[0];
155            }
156            for i in y1..=y2 {
157                *img.uget_mut((i as usize, x1 as usize, 0)) = draw_color[0];
158                *img.uget_mut((i as usize, x2 as usize, 0)) = draw_color[0];
159            }
160        }
161    }
162}
adder_codec_rs/utils/viz.rs

adder_codec_rs/utils/
viz.rs
