adder_codec_rs/utils/
viz.rs

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

adder_codec_rs/utils/
viz.rs
