diff --git a/learned.json b/learned.json index 23e4134..45ddc97 100644 --- a/learned.json +++ b/learned.json @@ -1,8 +1,4 @@ { - "learned_images": { - "/////8/wz/DP8Ofwx/OH8wP5A/kH+Z/5//////////8=": "47", - "//////////////////////////////////////////8=": "" - }, "learned_tracks": { "//////////+f/x//H/7f/N/A34HPn8+DD8AfwP////8=": "Magdalena Club", "/////z/wv/+/zz/vf+Z/5P/IH9gP2c+cD8CfwP////8=": "Rennvoort", diff --git a/src/configs/learned.default.json b/src/configs/learned.default.json index 3563fb1..8ecd131 100644 --- a/src/configs/learned.default.json +++ b/src/configs/learned.default.json @@ -1,4 +1,3 @@ { - "learned_images": {}, "learned_tracks": {} } \ No newline at end of file diff --git a/src/ocr_db.rs b/src/ocr_db.rs index 58d40d1..c8d694b 100644 --- a/src/ocr_db.rs +++ b/src/ocr_db.rs @@ -1,7 +1,8 @@ -use std::{collections::HashMap, sync::{RwLock}}; +use std::{collections::HashMap, sync::RwLock}; use crate::{ - config::{load_config_or_make_default, save_json_config}, ocr, + config::{load_config_or_make_default, save_json_config}, + ocr, }; use anyhow::Result; @@ -60,7 +61,10 @@ impl OcrDatabase { } pub fn learn(&self, hash: &str, val: char) { - self.learned_chars.write().unwrap().push((ImageHash::from_base64(hash).unwrap(), val)); + self.learned_chars + .write() + .unwrap() + .push((ImageHash::from_base64(hash).unwrap(), val)); } pub fn learn_phrase(&self, hashes: &[ImageHash], phrase: &str) -> Result<()> { @@ -83,17 +87,43 @@ impl OcrDatabase { pub fn ocr_char(&self, hash: &ImageHash) -> Option { let state = self.learned_chars.read().unwrap(); - let (_, c) = state.iter().min_by_key(|(learned_hash, _)| hash.dist(learned_hash))?; + let (_, c) = state + .iter() + .min_by_key(|(learned_hash, _)| hash.dist(learned_hash))?; Some(*c) } pub fn ocr_hashes(&self, hashes: &[ImageHash]) -> String { - let buffer: String = hashes.iter().filter_map(|hash| self.ocr_char(hash)).collect(); - buffer.trim_end_matches(|c: char| !c.is_alphanumeric()).to_owned() + let buffer: String = hashes + .iter() + .filter_map(|hash| self.ocr_char(hash)) + .collect(); + buffer + .trim_end_matches(|c: char| !c.is_alphanumeric()) + .to_owned() } - + pub fn ocr_image(&self, image: &RgbImage) -> String { let hashes = ocr::compute_box_hashes(image); self.ocr_hashes(&hashes) } } + +#[test] +fn test_ocr() { + let raw: RawOcrDatabase = + serde_json::from_str(include_str!("configs/ocr.default.json")).unwrap(); + let db: OcrDatabase = (&raw).into(); + + let image = image::load_from_memory(include_bytes!("test_data/test-image-3.png")).unwrap().to_rgb8(); + assert_eq!(db.ocr_image(&image), "00:30.625"); + + let image = image::load_from_memory(include_bytes!("test_data/test-image-4.png")).unwrap().to_rgb8(); + assert_eq!(db.ocr_image(&image), "00:20.296"); + + let image = image::load_from_memory(include_bytes!("test_data/test-image-num-1.png")).unwrap().to_rgb8(); + assert_eq!(db.ocr_image(&image), "1"); + + let image = image::load_from_memory(include_bytes!("test_data/test-image-blank.png")).unwrap().to_rgb8(); + assert_eq!(db.ocr_image(&image), ""); +} diff --git a/src/test_data/test-image-4.png b/src/test_data/test-image-4.png new file mode 100644 index 0000000..1a28c21 Binary files /dev/null and b/src/test_data/test-image-4.png differ diff --git a/src/test_data/test-image-blank.png b/src/test_data/test-image-blank.png new file mode 100644 index 0000000..bdd95c6 Binary files /dev/null and b/src/test_data/test-image-blank.png differ diff --git a/src/test_data/test-image-num-1.png b/src/test_data/test-image-num-1.png new file mode 100644 index 0000000..892eb73 Binary files /dev/null and b/src/test_data/test-image-num-1.png differ