aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
authorMelody Horn <melody@boringcactus.com>2021-12-03 19:33:19 -0700
committerMelody Horn <melody@boringcactus.com>2021-12-03 19:33:19 -0700
commit7cc6b6414d30e95e763baadee2f493c8f69329ac (patch)
treec0315d4c9d46caee9a96ee681a0be0a8b9960b7b /src
downloadqueue-go-brrr-7cc6b6414d30e95e763baadee2f493c8f69329ac.tar.gz
queue-go-brrr-7cc6b6414d30e95e763baadee2f493c8f69329ac.zip
piss
Diffstat (limited to 'src')
-rw-r--r--src/lib.rs158
-rw-r--r--src/ocr.rs50
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()
+}