local OCR
This commit is contained in:
Scott Pruett
parent
d3dfa83f32
commit
db2c73d0c1
|
|
@ -1 +1,2 @@
|
|||
/target
|
||||
/ocr_data
|
||||
|
|
@ -1761,6 +1761,12 @@ dependencies = [
|
|||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ppv-lite86"
|
||||
version = "0.2.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro-crate"
|
||||
version = "1.1.3"
|
||||
|
|
@ -1789,6 +1795,36 @@ dependencies = [
|
|||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"rand_chacha",
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_chacha"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
|
||||
dependencies = [
|
||||
"ppv-lite86",
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_core"
|
||||
version = "0.6.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
|
||||
dependencies = [
|
||||
"getrandom",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "raw-window-handle"
|
||||
version = "0.4.3"
|
||||
|
|
@ -2176,6 +2212,7 @@ dependencies = [
|
|||
"futures",
|
||||
"image 0.24.2",
|
||||
"img_hash",
|
||||
"rand",
|
||||
"reqwest",
|
||||
"scrap",
|
||||
"serde",
|
||||
|
|
|
|||
|
|
@ -25,3 +25,5 @@ img_hash = "3"
|
|||
|
||||
csv = "1"
|
||||
time = { version = "0.3", features = ["formatting", "local-offset"] }
|
||||
|
||||
rand = "0.8"
|
||||
|
|
@ -59,5 +59,7 @@
|
|||
"width": 30,
|
||||
"height": 30
|
||||
},
|
||||
"ocr_server_endpoint": "https://tesserver.spruett.dev/"
|
||||
"track_recognition_threshold": 10,
|
||||
"ocr_server_endpoint": "https://tesserver.spruett.dev/",
|
||||
"dump_frame_fraction": 0.05
|
||||
}
|
||||
48
learned.json
48
learned.json
|
|
@ -1,30 +1,34 @@
|
|||
{
|
||||
"learned_images": {
|
||||
"/////8/wz/DP8Ofwx/OH8wP5A/kH+Z/5//////////8=": "47"
|
||||
"/////8/wz/DP8Ofwx/OH8wP5A/kH+Z/5//////////8=": "47",
|
||||
"//////////////////////////////////////////8=": ""
|
||||
},
|
||||
"learned_tracks": {
|
||||
"/////z/wv+e/77/vf///7P/vP9iP38+YD4AfwP////8=": "Rennvoort",
|
||||
"//8//B/8z/0f/V/g/9w/+5/i3+Df/x/QH8D///////8=": "Tilksport GP",
|
||||
"///////////////5/+G/yZ/TH/IPyB/AH8D///////8=": "Tilksport Rallycross",
|
||||
"////5//Hf8I/4p/gz+DP99/zn4eflx+QP8D///////8=": "Whistle Valley",
|
||||
"//8//p/8z/0f/X/i/8w/25/C3+nPzx/EH8j///////8=": "Tilksport Club",
|
||||
"/////9//H//f/r/zP+QfzQ/YH8A/wj/I/+H///////8=": "Lost Lagoons",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efnx+CP8D///////8=": "Whistle Valley",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4eflx+EP8D///////8=": "Whistle Valley",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efn5+HP8D///////8=": "Whistle Valley",
|
||||
"//8f/J/8z/0f/X/iP8w/2Z/C3+nfzx/PH8n///////8=": "Tilksport GP",
|
||||
"///////////////5/8G/yZ/Dn+MPyB/OH8j/+f////8=": "Tilksport Rallycross",
|
||||
"/////z/wv+e/zz/vf+Z/7P/ID9gPkc+fD8Kf4P////8=": "Rennvoort",
|
||||
"/////3/wv+e/zz/vf+Z/5P/oH9iP2c+fD8KfwP////8=": "Rennvoort",
|
||||
"/////////////3/AH4APgA/QD8iP4R/4f/z///////8=": "Buffalo Hill",
|
||||
"//8f/o/+j/4f/z/8P/x//H/mf+7/9P/0//H/+f////8=": "Copperwood Club",
|
||||
"//8f/o/8z/6P/h/8P/1/8X/if+b/5P/0//H/+/////8=": "Copperwood Club",
|
||||
"////z//Af55/zj/jn/Hf/c/Bz48Pgw/A//////////8=": "Sugar Hill",
|
||||
"//////////+f/x//H/7f/N/A34HPn8+DD8AfwP////8=": "Magdalena Club",
|
||||
"/////z/wv/+/zz/vf+Z/5P/IH9gP2c+cD8CfwP////8=": "Rennvoort",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efnx+AP8D///////8=": "Whistle Valley",
|
||||
"//7//v/8f/x//U//H+Ifwj+ev9a/85/xH/j//P////8=": "Interstate",
|
||||
"//8//z//H/gfwB/ET9DP04/MH+D/4f////////////8=": "Maple Ridge",
|
||||
"//8/+B/yT/RPxR/J/9Ofxw/mz8DP2x/BH8D///////8=": "Thunder Point",
|
||||
"//8/5J/Ez8MfyX/g/8w/25/C3+nfzx/GH8D/+/////8=": "Tilksport GP"
|
||||
"/////9//H//f/r/zP+QfzQ/YH8A/wj/I/+H///////8=": "Lost Lagoons",
|
||||
"/////z/wv+e/77/vf///7P/vP9iP38+YD4AfwP////8=": "Rennvoort",
|
||||
"//8f/J/8z/0f/X/iP8w/2Z/C3+nfzx/PH8n///////8=": "Tilksport GP",
|
||||
"//7//v/8f/x//U//H+Ifwj+ev9a/85/xH/j//P////8=": "Interstate",
|
||||
"//8f/o/8z/6P/h/8P/1/8X/if+b/5P/0//H/+/////8=": "Copperwood Club",
|
||||
"////5//Hf8I/4p/gz+DP99/zn4eflx+QP8D///////8=": "Whistle Valley",
|
||||
"///////////////5/8G/yZ/Dn+MPyB/OH8j/+f////8=": "Tilksport Rallycross",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efnx+CP8D///////8=": "Whistle Valley",
|
||||
"//8//p/8z/0f/X/i/8w/25/C3+nPzx/EH8j///////8=": "Tilksport Club",
|
||||
"///////////////5/+G/yZ/TH/IPyB/AH8D///////8=": "Tilksport Rallycross",
|
||||
"//8/zo/Mz8afxh/iP+h/6H/of+b/9P/0//n/+/////8=": "Copperwood GP",
|
||||
"/////3/wv+e/zz/vf+Z/5P/oH9iP2c+fD8KfwP////8=": "Rennvoort",
|
||||
"////z//Af55/zj/jn/Hf/c/Bz48Pgw/A//////////8=": "Sugar Hill",
|
||||
"//8//B/8z/0f/V/g/9w/+5/i3+Df/x/QH8D///////8=": "Tilksport GP",
|
||||
"//8fzo/Mz8afwh/iP+B/6H/qf+b/5P/0//H/+/////8=": "Copperwood GP",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4eflx+EP8D///////8=": "Whistle Valley",
|
||||
"//8//z//H/gfwB/ET9DP04/MH+D/4f////////////8=": "Maple Ridge",
|
||||
"//8f/o/+j/4f/z/8P/x//H/mf+7/9P/0//H/+f////8=": "Copperwood Club",
|
||||
"/////z/wv+e/zz/vf+Z/7P/ID9gPkc+fD8Kf4P////8=": "Rennvoort",
|
||||
"//8/5J/Ez8MfyX/g/8w/25/C3+nfzx/GH8D/+/////8=": "Tilksport GP",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efn5+HP8D///////8=": "Whistle Valley",
|
||||
"////5//Hf+I/4p/gz+DP99/zn4efnx+AP8D///////8=": "Whistle Valley",
|
||||
"/////////////3/AH4APgA/QD8iP4R/4f/z///////8=": "Buffalo Hill"
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -157,3 +157,24 @@
|
|||
2022-06-02-02:28 (Copperwood Club),Copperwood Club,50s GT,9,23.071,22.037,100,56,74,
|
||||
2022-06-02-02:28 (Copperwood Club),Copperwood Club,50s GT,11,22.283,22.037,99,46,66,
|
||||
2022-06-02-02:28 (Copperwood Club),Copperwood Club,50s GT,12,22.544,22.037,99,36,57,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,1,34.263,32.543,94,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,2,33.221,32.543,91,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,3,33.492,32.543,20,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,4,32.543,32.543,20,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,5,33.534,32.543,20,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,6,32.971,32.543,20,,,
|
||||
2022-06-03-00:38 (Copperwood GP),Copperwood GP,Superlight,7,32.808,32.543,20,,,
|
||||
2022-06-03-00:37 (Copperwood GP),Copperwood GP,Superlight,1,32.513,32.513,,,,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,1,22.672,20.296,100,89,92,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,2,20.418,20.296,100,80,85,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,3,20.296,20.296,100,71,78,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,4,20.791,20.296,100,62,70,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,5,20.495,20.296,100,52,62,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,6,20.443,20.296,100,43,54,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,7,23.803,20.296,100,35,45,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,8,29.227,20.296,100,76,93,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,9,20.561,20.296,100,67,85,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,10,20.420,20.296,100,57,77,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,11,21.062,20.296,100,48,69,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,12,20.942,20.296,100,39,61,
|
||||
2022-06-03-00:30 (Magdalena Club),Magdalena Club,60s GP,13,20.576,20.296,100,30,52,
|
||||
|
|
|
|||
|
|
|
@ -167,15 +167,16 @@ fn add_saved_frame(
|
|||
fn run_loop_once(capturer: &mut Capturer, state: &SharedAppState) -> Result<()> {
|
||||
let frame = capture::get_frame(capturer)?;
|
||||
|
||||
let (config, learned_config, ocr_cache) = {
|
||||
let (config, learned_config, ocr_cache, should_sample) = {
|
||||
let locked = state.lock().unwrap();
|
||||
(
|
||||
locked.config.clone(),
|
||||
locked.learned.clone(),
|
||||
locked.ocr_cache.clone(),
|
||||
locked.should_sample_ocr_data
|
||||
)
|
||||
};
|
||||
let ocr_results = ocr::ocr_all_regions(&frame, config.clone(), learned_config, ocr_cache);
|
||||
let ocr_results = ocr::ocr_all_regions(&frame, config.clone(), learned_config, ocr_cache, should_sample);
|
||||
|
||||
if state.lock().unwrap().debug_frames {
|
||||
let debug_frames = save_frames_from(&frame, config.as_ref(), &ocr_results);
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ pub struct Config {
|
|||
pub use_ocr_cache: Option<bool>,
|
||||
pub ocr_interval_ms: Option<u64>,
|
||||
pub track_recognition_threshold: Option<u32>,
|
||||
pub dump_frame_fraction: Option<f64>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use image::{RgbImage, DynamicImage, Rgb};
|
||||
use image::{DynamicImage, Rgb, RgbImage};
|
||||
use img_hash::{image::GenericImageView, ImageHash};
|
||||
|
||||
use crate::image_processing;
|
||||
|
|
@ -8,7 +8,7 @@ struct BoundingBox {
|
|||
x: u32,
|
||||
y: u32,
|
||||
width: u32,
|
||||
height: u32
|
||||
height: u32,
|
||||
}
|
||||
|
||||
fn column_has_any_dark(image: &RgbImage, x: u32) -> bool {
|
||||
|
|
@ -25,13 +25,13 @@ fn row_has_any_dark(image: &RgbImage, y: u32, start_x: u32, width: u32) -> bool
|
|||
for x in start_x..(start_x + width) {
|
||||
let [r, g, b] = image.get_pixel(x, y).0;
|
||||
if r < 100 && g < 100 && b < 100 {
|
||||
return true
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn take_while<F: Fn(u32) -> bool>(image: &RgbImage, x: &mut u32, max: u32, f: F) {
|
||||
fn take_while<F: Fn(u32) -> bool>(x: &mut u32, max: u32, f: F) {
|
||||
while *x < max && f(*x) {
|
||||
*x = *x + 1;
|
||||
}
|
||||
|
|
@ -41,43 +41,69 @@ fn get_character_bounding_boxes(image: &RgbImage) -> Vec<BoundingBox> {
|
|||
let mut x = 0;
|
||||
let mut boxes = Vec::new();
|
||||
while x < image.width() {
|
||||
take_while(image, &mut x, image.width(), |x| !column_has_any_dark(image, x));
|
||||
take_while(&mut x, image.width(), |x| !column_has_any_dark(image, x));
|
||||
|
||||
let start_x = x;
|
||||
take_while(image, &mut x, image.width(), |x| column_has_any_dark(image, x));
|
||||
take_while(&mut x, image.width(), |x| column_has_any_dark(image, x));
|
||||
let width = x - start_x;
|
||||
|
||||
if width > 2 {
|
||||
if width >= 1 {
|
||||
let mut y = 0;
|
||||
take_while(image, &mut y, image.height(), |y| !row_has_any_dark(image, y, start_x, width));
|
||||
take_while(&mut y, image.height(), |y| {
|
||||
!row_has_any_dark(image, y, start_x, width)
|
||||
});
|
||||
|
||||
let start_y = y;
|
||||
|
||||
let mut inverse_y = 1;
|
||||
take_while(image, &mut inverse_y, image.height(), |y| !row_has_any_dark(image, image.height() - y, start_x, width));
|
||||
let end_y = image.height() - inverse_y - 1;
|
||||
let mut inverse_y = 0;
|
||||
take_while(&mut inverse_y, image.height(), |y| {
|
||||
!row_has_any_dark(image, image.height() - 1 - y, start_x, width)
|
||||
});
|
||||
let end_y = image.height() - inverse_y;
|
||||
let height = end_y - start_y;
|
||||
if height >= 1 {
|
||||
boxes.push(BoundingBox {
|
||||
x,
|
||||
x: start_x,
|
||||
y: start_y,
|
||||
width,
|
||||
height: end_y - start_y,
|
||||
height,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
boxes
|
||||
}
|
||||
|
||||
fn trim_to_bounding_box(image: &RgbImage, bounding_box: &BoundingBox) -> RgbImage {
|
||||
let mut buffer = RgbImage::new(bounding_box.width, bounding_box.height);
|
||||
const PADDING: u32 = 2;
|
||||
let mut buffer = RgbImage::from_pixel(
|
||||
bounding_box.width + 2 * PADDING,
|
||||
bounding_box.height + 2 * PADDING,
|
||||
Rgb([0xFF, 0xFF, 0xFF]),
|
||||
);
|
||||
for y in 0..bounding_box.height {
|
||||
for x in 0..bounding_box.width {
|
||||
buffer.put_pixel(x, y, *image.get_pixel(bounding_box.x + x, bounding_box.y + y));
|
||||
buffer.put_pixel(
|
||||
x + PADDING,
|
||||
y + PADDING,
|
||||
*image.get_pixel(bounding_box.x + x, bounding_box.y + y),
|
||||
);
|
||||
}
|
||||
}
|
||||
buffer
|
||||
}
|
||||
|
||||
fn compute_box_hashes(image: &RgbImage) -> Vec<String> {
|
||||
pub fn bounding_box_images(image: &RgbImage) -> Vec<RgbImage> {
|
||||
let mut trimmed = Vec::new();
|
||||
|
||||
let boxes = get_character_bounding_boxes(image);
|
||||
for bounding_box in boxes {
|
||||
trimmed.push(trim_to_bounding_box(image, &bounding_box));
|
||||
}
|
||||
trimmed
|
||||
}
|
||||
|
||||
pub fn compute_box_hashes(image: &RgbImage) -> Vec<String> {
|
||||
let mut hashes = Vec::new();
|
||||
|
||||
let boxes = get_character_bounding_boxes(image);
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ mod ocr;
|
|||
mod state;
|
||||
mod stats_writer;
|
||||
mod local_ocr;
|
||||
mod training_ui;
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
|
|
@ -30,6 +31,10 @@ use state::{AppState, DebugOcrFrame, LapState, OcrCache, RaceState, SharedAppSta
|
|||
use stats_writer::export_race_stats;
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
let mode = std::env::args().nth(1).unwrap_or_default().to_string();
|
||||
if mode == "train" {
|
||||
return training_ui::training_ui();
|
||||
}
|
||||
let app_state = AppState {
|
||||
config: Arc::new(Config::load().unwrap()),
|
||||
learned: Arc::new(LearnedConfig::load().unwrap()),
|
||||
|
|
@ -299,7 +304,7 @@ fn open_debug_lap(
|
|||
) {
|
||||
if let Some(screenshot_bytes) = &lap.screenshot {
|
||||
let screenshot = from_png_bytes(screenshot_bytes);
|
||||
let ocr_results = ocr::ocr_all_regions(&screenshot, config.clone(), learned, ocr_cache);
|
||||
let ocr_results = ocr::ocr_all_regions(&screenshot, config.clone(), learned, ocr_cache, false);
|
||||
let debug_lap = DebugLap {
|
||||
screenshot: RetainedImage::from_image_bytes("debug-lap", &to_png_bytes(&screenshot))
|
||||
.unwrap(),
|
||||
|
|
@ -397,6 +402,7 @@ impl eframe::App for AppUi {
|
|||
|
||||
ui.separator();
|
||||
ui.checkbox(&mut state.debug_frames, "Debug OCR regions");
|
||||
ui.checkbox(&mut state.should_sample_ocr_data, "Dump OCR training frames");
|
||||
});
|
||||
egui::CentralPanel::default().show(ctx, |ui| {
|
||||
egui::ScrollArea::vertical().show(ui, |ui| {
|
||||
|
|
|
|||
17
src/ocr.rs
17
src/ocr.rs
|
|
@ -54,7 +54,7 @@ async fn run_ocr_cached(
|
|||
hash: String,
|
||||
region: &crate::image_processing::Region,
|
||||
config: Arc<Config>,
|
||||
filtered_image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>>,
|
||||
filtered_image: &image::ImageBuffer<image::Rgb<u8>, Vec<u8>>,
|
||||
) -> Option<String> {
|
||||
let cached = {
|
||||
let locked = ocr_cache.read().unwrap();
|
||||
|
|
@ -66,7 +66,7 @@ async fn run_ocr_cached(
|
|||
return cached;
|
||||
}
|
||||
}
|
||||
match run_ocr(&filtered_image, &config.ocr_server_endpoint).await {
|
||||
match run_ocr(filtered_image, &config.ocr_server_endpoint).await {
|
||||
Ok(v) => {
|
||||
if use_cache {
|
||||
ocr_cache.write().unwrap().insert(hash.clone(), v.clone());
|
||||
|
|
@ -83,6 +83,7 @@ pub async fn ocr_all_regions(
|
|||
config: Arc<Config>,
|
||||
learned: Arc<LearnedConfig>,
|
||||
ocr_cache: Arc<OcrCache>,
|
||||
should_sample: bool
|
||||
) -> HashMap<String, Option<String>> {
|
||||
let results = Arc::new(Mutex::new(HashMap::new()));
|
||||
|
||||
|
|
@ -100,8 +101,18 @@ pub async fn ocr_all_regions(
|
|||
let value = if let Some(learned_value) = learned.learned_images.get(&hash) {
|
||||
Some(learned_value.clone())
|
||||
} else {
|
||||
run_ocr_cached(ocr_cache, hash, ®ion, config, filtered_image).await
|
||||
run_ocr_cached(ocr_cache, hash, ®ion, config.clone(), &filtered_image).await
|
||||
};
|
||||
|
||||
if let Some(sample_fraction) = &config.dump_frame_fraction {
|
||||
if rand::random::<f64>() < *sample_fraction {
|
||||
let file_id = rand::random::<usize>();
|
||||
let img_filename = format!("ocr_data/{}.png", file_id);
|
||||
filtered_image.save(img_filename).unwrap();
|
||||
let value_filename = format!("ocr_data/{}.txt", file_id);
|
||||
std::fs::write(value_filename, value.clone().unwrap_or_default()).unwrap();
|
||||
}
|
||||
}
|
||||
results.lock().unwrap().insert(region.name, value);
|
||||
}));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -140,6 +140,7 @@ pub struct AppState {
|
|||
|
||||
pub debug_frames: bool,
|
||||
pub saved_frames: HashMap<String, DebugOcrFrame>,
|
||||
pub should_sample_ocr_data: bool,
|
||||
|
||||
pub config: Arc<Config>,
|
||||
pub learned: Arc<LearnedConfig>,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,207 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
io::Write,
|
||||
path::{PathBuf, Path},
|
||||
sync::{Arc, Mutex},
|
||||
thread,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use eframe::{
|
||||
egui::{self, Ui, Visuals},
|
||||
emath::Vec2,
|
||||
epaint::Color32,
|
||||
};
|
||||
use egui_extras::RetainedImage;
|
||||
use image::RgbImage;
|
||||
|
||||
use crate::{image_processing::to_png_bytes, local_ocr};
|
||||
|
||||
#[derive(Default)]
|
||||
struct TrainingUi {
|
||||
training_images: Vec<TrainingImage>,
|
||||
|
||||
current_image_index: usize,
|
||||
|
||||
learned_char_hashes: Vec<(String, char)>,
|
||||
}
|
||||
|
||||
struct TrainingImage {
|
||||
img_file: PathBuf,
|
||||
data_file: PathBuf,
|
||||
image: RgbImage,
|
||||
text: String,
|
||||
|
||||
ui_image: RetainedImage,
|
||||
char_images: Vec<RetainedImage>,
|
||||
char_hashes: Vec<String>,
|
||||
}
|
||||
|
||||
impl TrainingImage {
|
||||
fn save_ocr_text(&self) {
|
||||
std::fs::write(&self.data_file, &self.text).unwrap();
|
||||
}
|
||||
fn delete_data(&self) {
|
||||
let _ = std::fs::remove_file(&self.img_file);
|
||||
let _ = std::fs::remove_file(&self.data_file);
|
||||
}
|
||||
}
|
||||
|
||||
fn get_training_data_paths() -> Vec<(PathBuf, PathBuf)> {
|
||||
let mut data_paths = Vec::new();
|
||||
for path in std::fs::read_dir("ocr_data/").unwrap() {
|
||||
let path = path.unwrap();
|
||||
if let Some(ext) = path.path().extension() {
|
||||
if ext.to_string_lossy() == "png" {
|
||||
data_paths.push((path.path(), path.path().with_extension("txt")));
|
||||
}
|
||||
}
|
||||
}
|
||||
data_paths.sort();
|
||||
data_paths
|
||||
}
|
||||
|
||||
fn predict_ocr(hashes: &[(String, char)], hash: &str) -> Option<char> {
|
||||
let hash = img_hash::ImageHash::<Vec<u8>>::from_base64(hash).unwrap();
|
||||
let (_, best_char) = hashes.iter().min_by_key(|(learned_hash, c)| {
|
||||
img_hash::ImageHash::from_base64(learned_hash)
|
||||
.unwrap()
|
||||
.dist(&hash)
|
||||
})?;
|
||||
Some(*best_char)
|
||||
}
|
||||
|
||||
fn get_training_data() -> Vec<TrainingImage> {
|
||||
let mut data = Vec::new();
|
||||
for (img_file, ocr_file) in get_training_data_paths() {
|
||||
let buffer = std::fs::read(&img_file).unwrap();
|
||||
let image = image::load_from_memory(&buffer).unwrap().to_rgb8();
|
||||
let ocr_value = String::from_utf8(std::fs::read(&ocr_file).unwrap()).unwrap();
|
||||
let ui_image = load_retained_image(&image);
|
||||
|
||||
let char_images = local_ocr::bounding_box_images(&image)
|
||||
.iter()
|
||||
.map(load_retained_image)
|
||||
.collect();
|
||||
let char_hashes = local_ocr::compute_box_hashes(&image);
|
||||
data.push(TrainingImage {
|
||||
img_file,
|
||||
data_file: ocr_file,
|
||||
image,
|
||||
text: ocr_value,
|
||||
ui_image,
|
||||
char_images,
|
||||
char_hashes,
|
||||
});
|
||||
}
|
||||
data
|
||||
}
|
||||
|
||||
fn load_learned_hashes() -> Vec<(String, char)> {
|
||||
let path = Path::new("learned_chars.txt");
|
||||
if !path.exists() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let data = String::from_utf8(std::fs::read(path).unwrap()).unwrap();
|
||||
let mut parsed = Vec::new();
|
||||
for line in data.lines() {
|
||||
if let Some((c, hash)) = line.split_once(" ") {
|
||||
if let Some(c) = c.chars().nth(0) {
|
||||
parsed.push((hash.to_owned(), c));
|
||||
}
|
||||
}
|
||||
}
|
||||
parsed
|
||||
}
|
||||
|
||||
fn load_retained_image(image: &RgbImage) -> RetainedImage {
|
||||
RetainedImage::from_image_bytes("", &to_png_bytes(image)).unwrap()
|
||||
}
|
||||
|
||||
pub fn training_ui() -> anyhow::Result<()> {
|
||||
let options = eframe::NativeOptions::default();
|
||||
let state = TrainingUi {
|
||||
training_images: get_training_data(),
|
||||
learned_char_hashes: load_learned_hashes(),
|
||||
..Default::default()
|
||||
};
|
||||
eframe::run_native("OCR Trainer", options, Box::new(|_cc| Box::new(state)));
|
||||
}
|
||||
|
||||
impl eframe::App for TrainingUi {
|
||||
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
|
||||
ctx.set_visuals(Visuals::dark());
|
||||
egui::CentralPanel::default().show(ctx, |ui| {
|
||||
egui::ScrollArea::vertical().show(ui, |ui| {
|
||||
let current_image = &mut self.training_images[self.current_image_index];
|
||||
if ui.button("Skip").clicked() {
|
||||
self.current_image_index += 1;
|
||||
}
|
||||
if ui.button("Delete Data").clicked() {
|
||||
current_image.delete_data();
|
||||
self.current_image_index += 1;
|
||||
}
|
||||
if ui.button("Save OCR fix").clicked() {
|
||||
current_image.save_ocr_text();
|
||||
}
|
||||
if ui.button("Learn").clicked() {
|
||||
for (i, char) in current_image.text.chars().enumerate() {
|
||||
if let Some(hash) = current_image.char_hashes.get(i) {
|
||||
self.learned_char_hashes.push((hash.clone(), char));
|
||||
eprintln!("Learned {}={}", hash, char);
|
||||
}
|
||||
}
|
||||
self.current_image_index += 1;
|
||||
}
|
||||
if ui.button("Learn and delete").clicked() {
|
||||
for (i, char) in current_image.text.chars().enumerate() {
|
||||
if let Some(hash) = current_image.char_hashes.get(i) {
|
||||
self.learned_char_hashes.push((hash.clone(), char));
|
||||
eprintln!("Learned {}={}", hash, char);
|
||||
}
|
||||
}
|
||||
current_image.delete_data();
|
||||
self.current_image_index += 1;
|
||||
}
|
||||
if ui.button("Save learned results").clicked() {
|
||||
let mut buffer = String::new();
|
||||
for (hash, c) in &self.learned_char_hashes {
|
||||
buffer += &format!("{} {}\n", c, hash);
|
||||
}
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.append(true)
|
||||
.create(true)
|
||||
.write(true)
|
||||
.open("learned_chars.txt")
|
||||
.unwrap();
|
||||
file.write(buffer.as_bytes()).unwrap();
|
||||
}
|
||||
|
||||
current_image.ui_image.show(ui);
|
||||
ui.label("OCR value");
|
||||
ui.text_edit_singleline(&mut current_image.text);
|
||||
let predicted: String = current_image
|
||||
.char_hashes
|
||||
.iter()
|
||||
.filter_map(|hash| predict_ocr(&self.learned_char_hashes, hash))
|
||||
.collect();
|
||||
if predicted == current_image.text {
|
||||
ui.colored_label(Color32::GREEN, format!("Predicted: {}", predicted));
|
||||
} else {
|
||||
ui.colored_label(Color32::RED, format!("Predicted: {}", predicted));
|
||||
}
|
||||
ui.separator();
|
||||
for c in ¤t_image.char_images {
|
||||
c.show(ui);
|
||||
}
|
||||
for c in ¤t_image.char_hashes {
|
||||
ui.label(c);
|
||||
if let Some(predicted) = predict_ocr(&self.learned_char_hashes, &c) {
|
||||
ui.label(format!("Predicted: {}", predicted));
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue