Files
ros2_lane_detect/src/pipeline/highlight_extractor.rs
2025-05-11 16:59:38 +02:00

68 lines
2.4 KiB
Rust

use std::sync::{mpsc, Arc, Mutex};
use std::sync::mpsc::Receiver;
use std::thread;
use std::time::{Duration, Instant};
use rand::Rng;
use r2r::sensor_msgs::msg::Image;
pub type Pointcloud = Arc<Vec<Vec<f64>>>;
trait ExtractHighlights {
fn get_highlights(&self, threshold: u8, downsample: f64) -> Pointcloud;
}
impl ExtractHighlights for Image {
fn get_highlights(&self, threshold: u8, downsample: f64) -> Pointcloud {
let h_i = [
[-7.74914499e-03, 3.95733793e-18, 3.10353257e00],
[8.56519716e-18, 9.42313768e-05, -1.86052093e00],
[2.57498016e-18, -2.73825295e-03, 1.00000000e00],
];
let max_rng = (u32::MAX as f64 * downsample) as u32;
Arc::new(self.data
.iter()
.enumerate()
.filter(|(i, _)| *i as u32 / self.width >= self.height / 2)
.filter(|(_, pixel)| **pixel > threshold)
.filter(|_| {
let mut rng = rand::rng();
rng.random::<u32>() < max_rng
})
.map(|(i, _)| {
let x = i as f64 % self.width as f64;
let y = i as f64 / self.width as f64;
// Convert to homogeneous coordinates
let u = h_i[0][0] * x + h_i[0][1] * y + h_i[0][2];
let v = h_i[1][0] * x + h_i[1][1] * y + h_i[1][2];
let w = h_i[2][0] * x + h_i[2][1] * y + h_i[2][2];
let u_norm = u / w;
let v_norm = v / w;
vec![u_norm, v_norm]
})
//.filter(|point| point[0].abs() <= 5.0)
.collect())
}
}
pub fn create_extract_thread(image_rx: Receiver<Arc<Image>>, threshold: u8, downsample: f64) -> (Receiver<Pointcloud>, Arc<Mutex<(Duration, Duration)>>) {
let (tx, rx) = mpsc::sync_channel::<Pointcloud>(1);
let usage = Arc::new(Mutex::new((Duration::ZERO, Duration::ZERO)));
let local_usage = usage.clone();
thread::spawn(move || {
let mut t_loop = Instant::now();
for image in image_rx {
let start = Instant::now();
let poi = image.get_highlights(threshold, downsample);
let end = Instant::now();
tx.send(poi).unwrap();
*local_usage.lock().unwrap() = (end - start, t_loop.elapsed());
t_loop = Instant::now();
}
});
(rx, usage)
}