Refactor: Use state amchine

Author: gfgafn

examples/create_mask.rs

@@ -3,7 +3,7 @@
 
 use opencv::{
 	core::{bitwise_and, find_file, CommandLineParser, Point, Scalar, CV_8UC1, CV_8UC3},
-	highgui::{self, *},
+	highgui::{self, imshow},
 	imgcodecs::{imread, IMREAD_COLOR},
 	imgproc,
 	prelude::*,
@@ -19,6 +19,16 @@ use std::{
 
 const SOURCE_WINDOW: &str = "Source image";
 
+#[derive(Debug, Clone, Copy)]
+enum DrawingState {
+	Init,
+	DrawingMarkerPoint,
+	DrawingMarkerPointFinished,
+	DrawingMask,
+	DrawingMaskFinished,
+	Resetting,
+}
+
 fn main() {
 	let args: Vec<String> = env::args().collect();
 	let (argc, argv) = (args.len() as i32, args.iter().map(|s| s.as_str()).collect::<Vec<&str>>());
@@ -33,18 +43,23 @@ fn main() {
 	);
 
 	// HACK: replace with parser.get<String>("@input")
-	let input_image = "lena.jpg";
-	let input_image_path = find_file(input_image, true, false).expect("Cannot find input image");
-
-	let [src, mut img1, mut mask, mut final_img]: [Mat; 4];
+	let input_image = argv.into_iter().nth(2).unwrap_or("lena.jpg");
+	let input_image_path = find_file(input_image, true, false)
+		.map(|path| {
+			println!("find input_image {} in : {}", input_image, path);
+			path
+		})
+		.unwrap_or_else(|_| panic!("Cannot find input_image: {}", input_image));
+
+	let [src, mut next_frame, mut mask, mut final_img]: [Mat; 4];
 	src = imread(&input_image_path, IMREAD_COLOR).unwrap();
 	if src.empty() {
 		eprintln!("Error opening image: {}", input_image);
 		process::exit(-1);
 	}
 
-	highgui::named_window(SOURCE_WINDOW, WINDOW_AUTOSIZE).unwrap();
-	let mouse_event_data: (i32, i32, i32, i32) = Default::default();
+	highgui::named_window(SOURCE_WINDOW, highgui::WINDOW_AUTOSIZE).unwrap();
+	let mouse_event_data = (highgui::MouseEventTypes::EVENT_MOUSEWHEEL, 0, 0, 0);
 	let (mouse_event_data, should_handle_mouse_event) = (Arc::new(Mutex::new(mouse_event_data)), Arc::new(AtomicBool::new(false)));
 
 	let mouse_event_dispatcher = {
@@ -54,7 +69,7 @@ fn main() {
 		move |event: i32, x: i32, y: i32, flags: i32| {
 			// can intercept specific mouse events here to don't update the mouse_data
 			if let Ok(mut mouse_data) = mouse_data.lock() {
-				*mouse_data = (event, x, y, flags);
+				*mouse_data = (mouse_event_from_i32(event), x, y, flags);
 			}
 			should_handle_mouse_event.store(true, atomic::Ordering::Relaxed);
 		}
@@ -63,10 +78,9 @@ fn main() {
 
 	highgui::imshow(SOURCE_WINDOW, &src).unwrap();
 
-	let (mut num_points_has_add, mut drag_left_button, mut right_button_down) = (0, false, false);
-	let mut points = Vec::<Point>::new();
+	let (mut marker_points, mut drawing_state) = (Vec::<Point>::new(), DrawingState::Init);
 
-	img1 = Mat::new_size_with_default(src.size().unwrap(), CV_8UC3, Scalar::new(0.0, 0.0, 0.0, 0.0)).unwrap();
+	next_frame = Mat::zeros_size(src.size().unwrap(), CV_8UC3).unwrap().to_mat().unwrap();
 
 	loop {
 		// Press Esc to exit
@@ -75,104 +89,141 @@ fn main() {
 		}
 
 		let (mouse_event, x, y, _) = {
-			if !should_handle_mouse_event.load(atomic::Ordering::Relaxed) {
+			if should_handle_mouse_event.load(atomic::Ordering::Relaxed) {
 				continue;
-			}
-			should_handle_mouse_event.store(false, atomic::Ordering::Relaxed);
-
-			if let Ok(mouse_event_data) = mouse_event_data.lock() {
-				*mouse_event_data
 			} else {
-				continue;
-			}
-		};
+				should_handle_mouse_event.store(false, atomic::Ordering::Relaxed);
 
-		match mouse_event {
-			EVENT_LBUTTONDOWN => {
-				dbg!((mouse_event, x, y));
-
-				if !drag_left_button && !right_button_down {
-					if num_points_has_add == 0 {
-						img1 = src.clone();
-					}
-
-					let point = Point::new(x, y);
-					imgproc::circle(&mut img1, point, 2, Scalar::new(0., 0., 255., 0.), -1, imgproc::LINE_8, 0).unwrap();
-					points.push(point);
-					num_points_has_add += 1;
-					drag_left_button = true;
-
-					if num_points_has_add > 1 {
-						imgproc::line(
-							&mut img1,
-							points[num_points_has_add - 2],
-							point,
-							Scalar::new(0., 0., 255., 0.),
-							2,
-							imgproc::LINE_8,
-							0,
-						)
-						.unwrap();
-					}
-
-					imshow(SOURCE_WINDOW, &img1).unwrap();
+				if let Ok(mouse_event_data) = mouse_event_data.lock() {
+					*mouse_event_data
+				} else {
+					continue;
 				}
 			}
-			EVENT_LBUTTONUP => {
-				dbg!((mouse_event, x, y));
-
-				if drag_left_button {
-					imshow(SOURCE_WINDOW, &img1).unwrap();
-					drag_left_button = false;
-				}
-			}
-			EVENT_RBUTTONDOWN => {
-				dbg!((mouse_event, x, y));
+		};
 
-				right_button_down = true;
-				let mut img1 = src.clone();
+		drawing_state = self::state_transform(drawing_state, mouse_event);
 
-				if num_points_has_add > 0 {
-					let pts_mat = Mat::from_slice(points.as_slice()).unwrap();
-					imgproc::polylines(&mut img1, &pts_mat, true, Scalar::new(0., 0., 0., 0.), 2, imgproc::LINE_8, 0).unwrap();
+		match drawing_state {
+			DrawingState::Init | DrawingState::DrawingMarkerPointFinished => { /* do nothing */ }
+			DrawingState::DrawingMarkerPoint => {
+				if marker_points.is_empty() {
+					next_frame = src.clone();
 				}
 
-				imshow(SOURCE_WINDOW, &img1).unwrap();
-			}
-			EVENT_RBUTTONUP => {
-				dbg!((mouse_event, x, y));
-
-				right_button_down = false;
-				final_img = Mat::zeros_size(src.size().unwrap(), CV_8UC3).unwrap().to_mat().unwrap();
-
-				mask = Mat::zeros_size(src.size().unwrap(), CV_8UC1).unwrap().to_mat().unwrap();
-
-				let points_mat = Mat::from_slice(points.as_slice()).unwrap();
-				imgproc::fill_poly(
-					&mut mask,
-					&points_mat,
-					Scalar::new(255., 255., 255., 255.),
+				let point = Point::new(x, y);
+				imgproc::circle(
+					&mut next_frame,
+					point,
+					2,
+					Scalar::new(0., 0., 255., 0.),
+					-1,
 					imgproc::LINE_8,
 					0,
-					Point::default(),
 				)
 				.unwrap();
+				marker_points.push(point);
+
+				if marker_points.len() > 1 {
+					imgproc::line(
+						&mut next_frame,
+						marker_points[marker_points.len() - 2],
+						point,
+						Scalar::new(0., 0., 255., 0.),
+						2,
+						imgproc::LINE_8,
+						0,
+					)
+					.unwrap();
+				}
 
-				bitwise_and(&src, &src, &mut final_img, &mask).unwrap();
-
-				imshow("Mask", &mask).unwrap();
-				imshow("Result", &final_img).unwrap();
-				imshow(SOURCE_WINDOW, &img1).unwrap();
+				imshow(SOURCE_WINDOW, &next_frame).unwrap();
+			}
+			DrawingState::DrawingMask => {
+				if !marker_points.is_empty() {
+					next_frame = src.clone();
+
+					let pts_mat = Mat::from_slice(marker_points.as_slice()).unwrap();
+					imgproc::polylines(
+						&mut next_frame,
+						&pts_mat,
+						true,
+						Scalar::new(0., 0., 0., 0.),
+						2,
+						imgproc::LINE_8,
+						0,
+					)
+					.unwrap();
+
+					imshow(SOURCE_WINDOW, &next_frame).unwrap();
+				}
+			}
+			DrawingState::DrawingMaskFinished => {
+				if !marker_points.is_empty() {
+					final_img = Mat::zeros_size(src.size().unwrap(), CV_8UC3).unwrap().to_mat().unwrap();
+					mask = Mat::zeros_size(src.size().unwrap(), CV_8UC1).unwrap().to_mat().unwrap();
+
+					imgproc::fill_poly(
+						&mut mask,
+						&Mat::from_slice(marker_points.as_slice()).unwrap(),
+						Scalar::new(255., 255., 255., 255.),
+						imgproc::LINE_8,
+						0,
+						Point::default(),
+					)
+					.unwrap();
+
+					bitwise_and(&src, &src, &mut final_img, &mask).unwrap();
+
+					imshow("Mask", &mask).unwrap();
+					imshow("Result", &final_img).unwrap();
+					imshow(SOURCE_WINDOW, &next_frame).unwrap();
+				}
 			}
-			EVENT_MBUTTONDOWN => {
-				dbg!((mouse_event, x, y));
+			DrawingState::Resetting => {
+				if !marker_points.is_empty() {
+					marker_points.clear();
+					next_frame = src.clone();
 
-				num_points_has_add = 0;
-				points.clear();
-				img1 = src.clone();
-				imshow(SOURCE_WINDOW, &img1).unwrap();
+					imshow(SOURCE_WINDOW, &next_frame).unwrap();
+				}
 			}
-			_ => {}
+		}
+	}
+}
+
+/// Converts an `i32` to a `opencv::highgui::MouseEventTypes`
+///
+/// # Panics
+///
+/// Panics if the argument less than 0 or greater than 11.
+fn mouse_event_from_i32(value: i32) -> opencv::highgui::MouseEventTypes {
+	(value.gt(&(opencv::highgui::MouseEventTypes::EVENT_MOUSEHWHEEL as i32) /* 11 */)
+		|| (value.lt(&(opencv::highgui::MouseEventTypes::EVENT_MOUSEMOVE as i32) /* 0 */)))
+	.then(|| panic!("Invalid cv::highgui::MouseEventTypes value: {}", value));
+
+	// Safe because of the previous check
+	unsafe { std::mem::transmute(value) }
+}
+
+fn state_transform(drawing_state: DrawingState, mouse_event: highgui::MouseEventTypes) -> DrawingState {
+	use self::DrawingState::*;
+	use opencv::highgui::MouseEventTypes::*;
+
+	match (&drawing_state, mouse_event) {
+		(Init, EVENT_LBUTTONDOWN) => DrawingMarkerPoint,
+		(DrawingMarkerPoint, EVENT_LBUTTONUP) => DrawingMarkerPointFinished,
+		(DrawingMarkerPointFinished, EVENT_LBUTTONDOWN) => DrawingMarkerPoint,
+		(DrawingMarkerPointFinished, EVENT_RBUTTONDOWN) => DrawingMask,
+		(DrawingMask, EVENT_RBUTTONUP) => DrawingMaskFinished,
+		(Init | DrawingMarkerPointFinished | DrawingMaskFinished, EVENT_MBUTTONDOWN) => Resetting,
+		(Resetting, EVENT_MBUTTONUP) => Init,
+		_ => {
+			println!(
+				"Invalid state transition from {:?} with event {:?}",
+				drawing_state, mouse_event
+			);
+			drawing_state
 		}
 	}
 }