diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/lib.rs | 158 | ||||
| -rw-r--r-- | src/ocr.rs | 50 |
2 files changed, 208 insertions, 0 deletions
diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..8f28199 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,158 @@ +use image::{DynamicImage, ImageFormat}; +use std::rc::Rc; +use std::sync::atomic::Ordering::Relaxed; +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; + +use js_sys::{ArrayBuffer, Uint8Array}; +use wasm_bindgen::prelude::*; +use wasm_bindgen::JsCast; +use wasm_bindgen_futures::*; +use web_sys::*; + +mod ocr; + +// When the `wee_alloc` feature is enabled, use `wee_alloc` as the global +// allocator. +#[cfg(feature = "wee_alloc")] +#[global_allocator] +static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT; + +async fn blob_to_image(blob: Blob) -> DynamicImage { + let array_buffer_promise = blob.array_buffer(); + let array_buffer: ArrayBuffer = JsFuture::from(array_buffer_promise) + .await + .unwrap() + .dyn_into() + .unwrap(); + let image_array = Uint8Array::new(&array_buffer); + let image_data = image_array.to_vec(); + image::load_from_memory_with_format(&image_data, ImageFormat::Png).unwrap() +} + +fn image_to_data_uri(image: &DynamicImage) -> String { + let mut result = String::from("data:image/png;base64,"); + let mut enc = base64::write::EncoderStringWriter::from(&mut result, base64::STANDARD); + image.write_to(&mut enc, ImageFormat::Png).unwrap(); + let _ = enc.into_inner(); + result +} + +fn get_queue_size(image: DynamicImage) -> DynamicImage { + let useful_region = image.crop_imm(911, 514, 70, 20); + + useful_region +} + +#[cfg(test)] +#[test] +fn test_queue_size() { + let image = image::open("data/2021-12-03 18-22-11.png").unwrap(); + let queue_size = get_queue_size(image); + queue_size.save("data/test.png").unwrap(); + let result = ocr::ocr(queue_size); + assert_eq!(result, 4_480); +} + +async fn do_boot() { + let width = Rc::new(AtomicU32::new(0)); + let height = Rc::new(AtomicU32::new(0)); + + let streaming = AtomicBool::new(false); + + let window = window().unwrap(); + let document = window.document().unwrap(); + + let video: HtmlVideoElement = document + .get_element_by_id("video") + .unwrap() + .dyn_into() + .unwrap(); + let canvas: HtmlCanvasElement = document + .get_element_by_id("canvas") + .unwrap() + .dyn_into() + .unwrap(); + let photo = document.get_element_by_id("photo").unwrap(); + let start_button = document.get_element_by_id("startbutton").unwrap(); + + let navigator = window.navigator(); + let media_devices = navigator.media_devices().unwrap(); + let constraints = { + let mut constraints = DisplayMediaStreamConstraints::new(); + constraints.audio(&JsValue::FALSE); + constraints.video(&JsValue::TRUE); + constraints + }; + let capture_stream_promise = media_devices + .get_display_media_with_constraints(&constraints) + .unwrap(); + let capture_stream = JsFuture::from(capture_stream_promise).await.unwrap(); + let capture_stream: MediaStream = capture_stream.dyn_into().unwrap(); + video.set_src_object(Some(&capture_stream)); + let video_listener = { + let width = width.clone(); + let height = height.clone(); + let video = video.clone(); + let also_video = video.clone(); + let canvas = canvas.clone(); + gloo::events::EventListener::new(&video, "canplay", move |_event| { + let _ = also_video.play().unwrap(); + if !streaming.load(Relaxed) { + width.store(also_video.video_width(), Relaxed); + height.store(also_video.video_height(), Relaxed); + canvas.set_attribute("width", &format!("{}", width.load(Relaxed))); + canvas.set_attribute("height", &format!("{}", height.load(Relaxed))); + streaming.store(true, Relaxed); + } + }) + }; + video_listener.forget(); + + let start_button_event_listener = { + let photo = photo.clone(); + gloo::events::EventListener::new(&start_button, "click", move |_event| { + let context = canvas.get_context("2d").unwrap().unwrap(); + let context: CanvasRenderingContext2d = context.dyn_into().unwrap(); + canvas.set_width(width.load(Ordering::Relaxed)); + canvas.set_height(height.load(Ordering::Relaxed)); + context + .draw_image_with_html_video_element_and_dw_and_dh( + &video, + 0.0, + 0.0, + width.load(Relaxed).into(), + height.load(Relaxed).into(), + ) + .unwrap(); + + canvas + .to_blob( + Closure::once_into_js(|data: Blob| { + console::log_1(&JsValue::from("hewwo?")); + // lifetime bullshit here + let window = web_sys::window().unwrap(); + let document = window.document().unwrap(); + let photo = document.get_element_by_id("photo").unwrap(); + spawn_local(async move { + let image = blob_to_image(data).await; + let useful_region = get_queue_size(image); + let data_uri = image_to_data_uri(&useful_region); + photo.set_attribute("src", &data_uri).unwrap(); + console::log_1(&JsValue::from("uhhhh hi")); + }) + }) + .dyn_ref() + .unwrap(), + ) + .unwrap(); + }) + }; + start_button_event_listener.forget(); +} + +#[wasm_bindgen] +pub fn boot() { + console_error_panic_hook::set_once(); + + spawn_local(do_boot()); +} diff --git a/src/ocr.rs b/src/ocr.rs new file mode 100644 index 0000000..bc7e535 --- /dev/null +++ b/src/ocr.rs @@ -0,0 +1,50 @@ +use image::{DynamicImage, GrayImage, ImageFormat}; +use imageproc::template_matching::MatchTemplateMethod; + +const REFERENCE_PNGS: [&[u8]; 10] = [ + include_bytes!("../data/0.png"), + include_bytes!("../data/1.png"), + include_bytes!("../data/2.png"), + include_bytes!("../data/3.png"), + include_bytes!("../data/4.png"), + include_bytes!("../data/5.png"), + include_bytes!("../data/6.png"), + include_bytes!("../data/7.png"), + include_bytes!("../data/8.png"), + include_bytes!("../data/9.png"), +]; + +fn x_matches(image: &GrayImage, template: &GrayImage) -> Vec<u32> { + let match_values = imageproc::template_matching::match_template( + image, + template, + MatchTemplateMethod::CrossCorrelationNormalized, + ); + match_values + .enumerate_pixels() + .filter(|(_x, _y, pix)| pix.0[0] > 0.9) + .map(|(x, _y, _pix)| x) + .collect() +} + +pub fn ocr(image: DynamicImage) -> u32 { + let grayscale_image = image::imageops::grayscale(&image); + let mut digit_x_positions: Vec<(u8, u32)> = (0..10) + .flat_map(|i| { + let png_i = + image::load_from_memory_with_format(REFERENCE_PNGS[i as usize], ImageFormat::Png) + .unwrap(); + let grey_png_i = image::imageops::grayscale(&png_i); + x_matches(&grayscale_image, &grey_png_i) + .into_iter() + .map(move |x| (i, x)) + }) + .collect(); + digit_x_positions.sort_by_key(|&(_i, x)| x); + let digits: String = digit_x_positions + .into_iter() + .map(|(i, _x)| format!("{}", i)) + .collect(); + dbg!(&digits); + digits.parse().unwrap() +} |