1
Fork 0
cruisecontrol/simcontroller/src/main.rs
2024-09-10 21:04:53 -04:00

190 lines
7.7 KiB
Rust

use camera::Camera;
use opencv::{core::{bitwise_and, bitwise_and_def, bitwise_or_def, in_range, KeyPoint, KeyPointTraitConst, Mat, MatTraitConst, MatTraitConstManual, Ptr, Size, VecN, Vector, CV_64F, CV_8UC4}, features2d::{draw_matches_def, BFMatcher, ORB}, gapi::bgr2_luv, highgui::{self, MouseCallback}, imgcodecs::ImreadModes, imgproc::{self, cvt_color_def, gaussian_blur, gaussian_blur_def, sobel_def, COLOR_BGR2GRAY, COLOR_BGR2HLS}, prelude::{DescriptorMatcherTrait, DescriptorMatcherTraitConst, Feature2DTrait}};
use webots::prelude::*;
use std::{f64::consts::TAU, time::Duration};
mod camera;
const TIME_STEP: Duration = Duration::from_millis(20);
const MAX_SPEED: f64 = TAU;
// const tweaker only supports floating point, these are mapped
#[const_tweaker::tweak]
const LOWERB_H: f32 = 18. / 180.;
#[const_tweaker::tweak]
const LOWERB_L: f32 = 23. / 255.;
#[const_tweaker::tweak]
const LOWERB_S: f32 = 36. / 255.;
#[const_tweaker::tweak]
const UPPERB_H: f32 = 92. / 180.;
#[const_tweaker::tweak]
const UPPERB_L: f32 = 151. / 255.;
#[const_tweaker::tweak]
const UPPERB_S: f32 = 169. / 255.;
struct MyRobot {
camera: Camera,
left_motor: Motor,
right_motor: Motor,
window: &'static str,
orb_detector: Ptr<ORB>,
matcher: BFMatcher,
last_frame_descriptors: Option<Mat>,
last_frame_keypoints: Option<Vector<KeyPoint>>,
last_image: Option<Mat>,
direction: f32,
rotational_velocity: f32,
}
impl Robot for MyRobot {
fn time_step(&self) -> Duration {
TIME_STEP
}
fn init() -> Self {
let camera = Camera::new("camera7");
let left_motor = Motor::new("motorl");
let right_motor = Motor::new("motorr");
left_motor.set_position(f64::INFINITY);
right_motor.set_position(f64::INFINITY);
left_motor.set_velocity(0.1 * MAX_SPEED);
right_motor.set_velocity(0.1 * MAX_SPEED);
let window = "processed";
//highgui::named_window(window, 1).unwrap();
//highgui::set_mouse_callback(window, |x,y,_,_| {dbg!(x+y);}).unwrap();
let orb_detector = ORB::create_def().unwrap();
let matcher = opencv::features2d::BFMatcher::new_def().unwrap();
//let last_frame_descriptors = opencv::core::Mat::default();
Self {
camera,
left_motor,
right_motor,
window,
orb_detector,
matcher,
last_frame_descriptors: None,
last_frame_keypoints: None,
last_image: None,
direction: 0.,
rotational_velocity: 0.,
}
}
fn step(&mut self, _time: StepTime) {
self.camera.enable(20);
if let Ok(img) = self.camera.get_image() { // BGRA
let img = opencv::core::Mat::new_rows_cols_with_bytes::<VecN<u8,4>>(240, 320, &img).unwrap();
//opencv::imgcodecs::imwrite_def("cam.png", &img).unwrap();
//highgui::imshow(self.window, &img).unwrap();
let mut blurred = opencv::core::Mat::default();
let mut edge = opencv::core::Mat::default();
gaussian_blur_def(&img, &mut blurred, Size::new(3, 3), 0.).unwrap();
//sobel_def(&blurred, &mut edge, CV_64F, 0, 2).unwrap();
cvt_color_def(&blurred, &mut edge, COLOR_BGR2HLS).unwrap();
highgui::imshow(self.window, &blurred).unwrap();
let mut mask_a = opencv::core::Mat::default();
let mut mask_b = opencv::core::Mat::default();
let mut mask = opencv::core::Mat::default();
opencv::imgcodecs::imwrite_def("blurred.png", &blurred).unwrap();
println!("Lowrerh {}", LOWERB_H);
let lowerb = dbg!([(*LOWERB_H*255.) as u8,(*LOWERB_L*255.) as u8,(*LOWERB_S*255.) as u8,255]);
let upperb = [(*UPPERB_H*255.) as u8,(*UPPERB_L*255.) as u8,(*UPPERB_S*255.) as u8,255];
in_range(&blurred, &lowerb, &upperb, &mut mask_a).unwrap(); // h low
// 103/180
in_range(&blurred, &[0,30,4,255], &[52,209,63,255], &mut mask_b).unwrap(); // hh28
highgui::imshow("a", &mask_a).unwrap();
//highgui::imshow("b", &mask_b).unwrap();
dbg!(blurred.row(240/2).unwrap().col(320/2).unwrap().data_bytes().unwrap());
let mut range = opencv::core::Mat::default();
bitwise_or_def(&mask_a, &mask_a, &mut mask).unwrap();
bitwise_and(&img, &img, &mut range, &mask).unwrap();
highgui::imshow("edge", &range).unwrap();
let mask = opencv::core::Mat::default();
let mut keypoints = opencv::core::Vector::default();
let mut descriptors = opencv::core::Mat::default();
self.orb_detector.detect_and_compute_def(&blurred, &mask, &mut keypoints, &mut descriptors).unwrap();
let mut debugout = opencv::core::Mat::default();
opencv::features2d::draw_keypoints(&img, &keypoints, &mut debugout, opencv::core::VecN([0.,255.,255.,255.]), opencv::features2d::DrawMatchesFlags::DEFAULT).unwrap();
//highgui::imshow("keypoints", &debugout).unwrap();
let mut matches: Vector<opencv::core::DMatch> = Vector::new();
if let Some(mut last_frame_descriptors) = self.last_frame_descriptors.as_ref() {
if let Err(e) = self.matcher.train_match_def(&mut last_frame_descriptors, &mut descriptors, &mut matches) {
eprintln!("no features to match");
}
let matches = matches.iter().filter(|m| m.distance < 105.).collect();
draw_matches_def(self.last_image.as_ref().unwrap(), self.last_frame_keypoints.as_ref().unwrap(), &img , &keypoints, &matches, &mut debugout).unwrap();
highgui::imshow("matches", &debugout).unwrap();
let mut horiz = Vec::new();
for hit in matches {
if hit.distance > 45. {continue;}
let last = self.last_frame_keypoints.as_ref().unwrap().get(hit.query_idx as usize).unwrap();
let current = keypoints.get(hit.train_idx as usize).unwrap();
let dx = last.pt().x - current.pt().x;
horiz.push(dx);
}
if horiz.len() > 0 {
let avg = horiz.iter().sum::<f32>()/horiz.len() as f32;
let lerp = 0.1 + horiz.len() as f32 * 0.05;
let lerp = lerp.max(0.65);
self.rotational_velocity = self.rotational_velocity * (1.-lerp) + avg * lerp;
//self.direction += avg;
}
self.direction += self.rotational_velocity;
}
self.last_frame_descriptors = Some(descriptors);
self.last_frame_keypoints = Some(keypoints);
self.last_image = Some(img.clone_pointee());
}
highgui::poll_key().unwrap();
const DEG_TO_PX: f64 = 320. / 62.2;
let goal = if _time.elapsed.as_secs() % 8 < 4 {
90. * DEG_TO_PX
} else {
0. * DEG_TO_PX
};
let error = goal - self.direction as f64;
let mut error = error * 0.02;
if error.abs() < 0.5 {
error = 0.;
}
dbg!(self.direction);
dbg!(self.rotational_velocity);
let mut left_speed = error * MAX_SPEED;
let mut right_speed = -error * MAX_SPEED;
//self.left_motor.set_velocity(left_speed.clamp(-10., 10.));
//self.right_motor.set_velocity(right_speed.clamp(-10., 10.));
self.left_motor.set_velocity(0.);
self.right_motor.set_velocity(0.);
}
}
fn main() {
println!("Rust controller has started");
webots::run_robot::<MyRobot>();
}