diff --git a/.gitignore b/.gitignore index ea8c4bf..7977566 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ /target +*.rten +download-models.sh diff --git a/Cargo.lock b/Cargo.lock index 235cd07..8797797 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,7 +31,7 @@ checksum = "c980ee35e870bd1a4d2c8294d4c04d0499e67bca1e4b5cefcc693c2fa00caea9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -63,9 +63,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bumpalo" -version = "3.14.0" +version = "3.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec" +checksum = "d32a994c2b3ca201d9b263612a374263f05e7adde37c4707f693dcd375076d1f" [[package]] name = "cairo-rs" @@ -117,6 +117,37 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" + +[[package]] +name = "either" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" + [[package]] name = "equivalent" version = "1.0.1" @@ -133,6 +164,16 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "flatbuffers" +version = "22.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ae1bfc84d904f75e7ef6f8796b020c606a9e8e271e2004c0a74f7edeedba45f" +dependencies = [ + "bitflags", + "rustc_version", +] + [[package]] name = "flume" version = "0.10.14" @@ -208,7 +249,7 @@ checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -566,6 +607,12 @@ version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +[[package]] +name = "libm" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" + [[package]] name = "lock_api" version = "0.4.11" @@ -634,6 +681,19 @@ dependencies = [ "memchr", ] +[[package]] +name = "ocrs" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfd2f495d2171a0e9eb972c2a8568b2dca7d56777e6ec7e54f974f36b532c340" +dependencies = [ + "rayon", + "rten", + "rten-imageproc", + "rten-tensor", + "wasm-bindgen", +] + [[package]] name = "once_cell" version = "1.19.0" @@ -683,7 +743,7 @@ checksum = "266c042b60c9c76b8d53061e52b2e0d1116abc57cefc8c5cd671619a56ac3690" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -700,9 +760,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pkg-config" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2900ede94e305130c13ddd391e0ab7cbaeb783945ae07a279c268cb05109c6cb" +checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" [[package]] name = "proc-macro-crate" @@ -756,6 +816,26 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rayon" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7237101a77a10773db45d62004a272517633fbcc3df19d96455ede1122e051" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "relm4" version = "0.6.2" @@ -781,18 +861,48 @@ checksum = "9340e2553c0a184a80a0bfa1dcf73c47f3d48933aa6be90724b202f9fbd24735" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", +] + +[[package]] +name = "rten" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67842141586c527d25547bbd81ccfd60601940281a792bce14a3f72dc15a3fec" +dependencies = [ + "flatbuffers", + "libm", + "rayon", + "rten-tensor", + "rten-vecmath", + "smallvec", + "wasm-bindgen", +] + +[[package]] +name = "rten-imageproc" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9ceab28b8dcb8925ebc6e51f0038fd16e1cd68229e596eadef826ec7794747f" +dependencies = [ + "rten-tensor", ] [[package]] name = "rten-tensor" -version = "0.4.0" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b994b2c4597280249d6beb99fe01dd0adf06b0f002c2049a07708f09aed55c3a" +checksum = "b00a55ca04d3219957737f87a9fc3f6eee6770e1f1643c2094c032c90c43f0d2" dependencies = [ "smallvec", ] +[[package]] +name = "rten-vecmath" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae896b8b5bea024dab45f9c8ee71051dfbf2a20a83b0e4e48cea5c4e7096e84" + [[package]] name = "rustc-demangle" version = "0.1.23" @@ -837,7 +947,7 @@ checksum = "33c85360c95e7d137454dc81d9a4ed2b8efd8fbe19cee57357b32b9771fccb67" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -886,9 +996,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.48" +version = "2.0.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f3531638e407dfc0814761abb7c00a5b54992b849452a0646b7f65c9f770f3f" +checksum = "915aea9e586f80826ee59f8453c1101f9d1c4b3964cd2460185ee8e299ada496" dependencies = [ "proc-macro2", "quote", @@ -931,7 +1041,7 @@ checksum = "a953cb265bef375dae3de6663da4d3804eee9682ea80d8e2542529b73c531c81" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -954,7 +1064,7 @@ dependencies = [ "serde", "serde_spanned", "toml_datetime", - "toml_edit 0.22.5", + "toml_edit 0.22.6", ] [[package]] @@ -979,15 +1089,15 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.22.5" +version = "0.22.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99e68c159e8f5ba8a28c4eb7b0c0c190d77bb479047ca713270048145a9ad28a" +checksum = "2c1b5fd4128cc8d3e0cb74d4ed9a9cc7c7284becd4df68f5f940e1ad123606f6" dependencies = [ "indexmap", "serde", "serde_spanned", "toml_datetime", - "winnow 0.6.0", + "winnow 0.6.1", ] [[package]] @@ -1009,7 +1119,7 @@ checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", ] [[package]] @@ -1066,7 +1176,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", "wasm-bindgen-shared", ] @@ -1088,7 +1198,7 @@ checksum = "642f325be6301eb8107a83d12a8ac6c1e1c54345a7ef1a9261962dfefda09e66" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.49", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -1103,7 +1213,9 @@ checksum = "4f186bd2dcf04330886ce82d6f33dd75a7bfcf69ecf5763b89fcde53b6ac9838" name = "waywrite" version = "0.1.0" dependencies = [ + "ocrs", "relm4", + "rten", "rten-tensor", ] @@ -1140,9 +1252,9 @@ dependencies = [ [[package]] name = "winnow" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b1dbce9e90e5404c5a52ed82b1d13fc8cfbdad85033b6f57546ffd1265f8451" +checksum = "d90f4e0f530c4c69f62b80d839e9ef3855edc9cba471a160c4d692deed62b401" dependencies = [ "memchr", ] diff --git a/Cargo.toml b/Cargo.toml index 7a99ab1..4cc6c5c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,5 +6,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] -relm4="0.6.*" -rten-tensor = "0.4.*" +relm4 = "0.6.*" +rten-tensor = "0.3.*" +rten = "0.3.1" +ocrs = "0.4.*" diff --git a/process_point.rs b/process_point.rs deleted file mode 100644 index e69de29..0000000 diff --git a/src/main.rs b/src/main.rs index 4548b74..abbcb6e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,7 @@ use gtk::prelude::{BoxExt, ButtonExt, GtkWindowExt, OrientableExt, WidgetExt, Ge use gtk::cairo::{Context, Operator}; use relm4::drawing::DrawHandler; -use waywrite::process_point::to_matrix; +use waywrite::process_point::print_written; use waywrite::Point; #[derive(Debug)] @@ -113,12 +113,7 @@ impl SimpleComponent for AppModel { match message { AppInput::Input => { - println!("\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n"); - let mat = to_matrix(&self.points); - for line in mat.iter() { - line.iter().for_each(|x| print!("{}", x)); - println!(); - } + print_written(&self.points).unwrap(); } AppInput::AddPoint((x, y)) => { self.points.push(Point { x, y, new_line: false }) @@ -144,8 +139,8 @@ fn draw(cx: &Context, points: &[Point]) { if !point.new_line { let last_point = &points[i - 1]; - cx.move_to(last_point.x, last_point.y); - cx.line_to(point.x, point.y); + cx.move_to(last_point.x as f64, last_point.y as f64); + cx.line_to(point.x as f64, point.y as f64); cx.set_source_rgb(200.0, 200.0, 200.0); cx.set_line_width(2.0); cx.stroke().expect("Failed to draw line"); diff --git a/src/process_point.rs b/src/process_point.rs index 4731772..665bda2 100644 --- a/src/process_point.rs +++ b/src/process_point.rs @@ -1,53 +1,171 @@ -use rten_tensor::NdTensor; +use std::fs; +use std::error::Error; +use std::time::Instant; + +use rten_tensor::{NdTensor, AsView}; +use rten::Model; +use ocrs::{OcrEngine, OcrEngineParams}; use crate::Point; -fn process(points: Vec) -> NdTensor { - - - NdTensor::zeros([1, 1, 1]) +const MATRIX_SIZE: usize = 800; + +pub fn print_written(points: &Vec) -> Result<(), Box> { + + let begin = Instant::now(); + let processed_data = process(points); + + println!("{:#?}", begin.elapsed()); + let begin = Instant::now(); + ocr(processed_data)?; + println!("{:#?}", begin.elapsed()); + + Ok(()) } -pub fn to_matrix(points: &Vec) -> Vec>{ +fn process(points: &Vec) -> NdTensor { - const MATRIX_SIZE: f64 = 200.0; + let matrix = to_matrix(points); + let y_len = matrix[0].len(); + let x_len = matrix[0][0].len(); + let data: Vec = matrix.into_iter().flatten().flatten().map(|f| f as f32).collect(); + + + NdTensor::from_data([1, y_len, x_len], data) +} + +fn ocr(data: NdTensor) -> Result<(), Box> { + + let detection_model_data = fs::read("text-detection.rten")?; + let rec_model_data = fs::read("text-recognition.rten")?; + + let detection_model = Model::load(&detection_model_data)?; + let rec_model = Model::load(&rec_model_data)?; + + + let ocr_engine = OcrEngine::new(OcrEngineParams { + detection_model: Some(detection_model), + recognition_model: Some(rec_model), + ..Default::default() + })?; + + let input = ocr_engine.prepare_input(data.view())?; + + let word_rects = ocr_engine.detect_words(&input)?; + + let line_rects = ocr_engine.find_text_lines(&input, &word_rects); + + let line_texts = ocr_engine.recognize_text(&input, &line_rects)?; + + for line in line_texts + .iter() + .flatten() + .filter(|l| l.to_string().len() > 1) + { + println!("{}", line); + } + + + + Ok(()) + + +} + +fn line(x: f64, point1: (f64, f64), point2: (f64, f64)) -> f64 { + let slope = (point2.1 - point1.1) / (point2.0 - point1.0); + + let point = slope * (x - point1.0) + point1.1; + + point +} + +fn to_matrix(points: &Vec) -> Vec>> { let min_x = points.iter().min_by_key(|p| p.x as i32).unwrap().x; let min_y = points.iter().min_by_key(|p| p.y as i32).unwrap().y; let max_x = points.iter().max_by_key(|p| p.x as i32).unwrap().x; let max_y = points.iter().max_by_key(|p| p.y as i32).unwrap().y; - let x_len = max_x - min_x; let y_len = max_y - min_y; + let y_ratio = y_len / x_len; - let x_scale = (0.8 * MATRIX_SIZE) / x_len; - let y_scale = ((0.8 * MATRIX_SIZE) * y_ratio) / y_len; - let scaled_points = points.iter().map(|point| { - let x_scaled = ((point.x - min_x) * x_scale) + (0.1 * MATRIX_SIZE); - let y_scaled = ((point.y - min_y) * y_scale) + ((0.1 * MATRIX_SIZE) * y_ratio); + let x_size = MATRIX_SIZE as f64 * 0.5; + let y_size = (MATRIX_SIZE as f64 * 0.5) * y_ratio; - ((x_scaled, y_scaled), point.new_line) - }); + let x_offset = (MATRIX_SIZE as f64 - x_size) / 2.0; + let y_offset = ((MATRIX_SIZE as f64 * y_ratio) - y_size) / 2.0; - let mut matrix: Vec> = vec![Box::new([0.0; MATRIX_SIZE as u32 as usize]); (MATRIX_SIZE * y_ratio) as u64 as usize]; + let x_scale = x_size / x_len; + let y_scale = y_size / y_len; + + let mut matrix: Vec> = vec![vec![0.0; MATRIX_SIZE]; (MATRIX_SIZE as f64 * y_ratio) as usize]; + + + let scaled_points = points + .iter() + .map(|point| { + let x_scaled = ((point.x - min_x) * x_scale) + x_offset; + let y_scaled = ((point.y - min_y) * y_scale) + y_offset; + + ((x_scaled, y_scaled), point.new_line) + }).collect::>(); + + let line_width_x = MATRIX_SIZE as f64 / 80.0; + let line_width_y = (MATRIX_SIZE as f64 * y_ratio) / 80.0; + + let mut last_x = 0.0; + let mut last_y = 0.0; for ((x, y), newline) in scaled_points { - let start_x = x - (MATRIX_SIZE / 100.0); - let end_x = x + (MATRIX_SIZE / 100.0); - let start_y = y - ((MATRIX_SIZE / 100.0) * y_ratio); - let end_y = y + ((MATRIX_SIZE / 100.0) * y_ratio); - matrix.iter_mut().enumerate().for_each(|(mat_y, line)| line.iter_mut().enumerate().for_each(|(mat_x, val)| { - if (start_x < (mat_x as f64) && (mat_x as f64) < end_x) && (start_y < (mat_y as f64) && (mat_y as f64) < end_y) { - *val = 1.0; + if !newline { + let curr_x_start = x - (line_width_x / 2.0); + let curr_x_end = x + (line_width_x / 2.0); + + let last_x_start = last_x - (line_width_x / 2.0); + let last_x_end = last_x + (line_width_x / 2.0); + + let top_y: f64; + let bottom_y: f64; + + if y > last_y { + top_y = y + ((line_width_y / 2.0) * y_scale); + bottom_y = last_y - ((line_width_y / 2.0) * y_scale); + } else { + top_y = last_y + ((line_width_y / 2.0) * y_scale); + bottom_y = y - ((line_width_y / 2.0) * y_scale); } - })); + + let start_x = (last_x_start.min(curr_x_start)) as usize; + let end_x = (last_x_end.max(curr_x_end)) as usize + 1; + + + for x in start_x..(end_x + 1) { + + let left_line_y = line(x as f64, (last_x_start, last_y), (curr_x_start, y)); + let right_line_y = line(x as f64, (last_x_end, last_y), (curr_x_end, y)); + + let top_line = left_line_y.max(right_line_y); + let bottom_line = left_line_y.min(right_line_y); + + let top_line = top_line.min(top_y) as usize + 1; + let bottom_line = bottom_line.max(bottom_y) as usize; + + for line_y in bottom_line..(top_line + 1) { + matrix[line_y][x] = 1.0; + } + + } + } + + last_x = x; + last_y = y; + } - - - matrix + vec![matrix] }