1
Fork 0
shimmeringmoon/src/recognition/hyperglass.rs
prescientmoon 5796bd6865
Fix warning
Signed-off-by: prescientmoon <git@moonythm.dev>
2024-08-11 03:15:44 +02:00

291 lines
7 KiB
Rust

use freetype::Face;
use image::{DynamicImage, ImageBuffer, Luma};
use imageproc::{
contrast::{threshold, ThresholdType},
region_labelling::{connected_components, Connectivity},
};
use num::traits::Euclid;
use crate::{
bitmap::{Align, BitmapCanvas, Color, TextStyle},
context::Error,
logs::{debug_image_buffer_log, debug_image_log},
};
///! Hyperglass my own specialized OCR system
// {{{ ConponentVec
/// How many sub-segments to split each side into
const SPLIT_FACTOR: u32 = 5;
const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR) as usize;
#[derive(Debug, Clone)]
struct ComponentVec {
chunks: [f32; IMAGE_VEC_DIM],
}
impl ComponentVec {
// {{{ (Component => vector) encoding
fn from_component(components: &ComponentsWithBounds, component: u32) -> Result<Self, Error> {
let mut chunks = [0.0; IMAGE_VEC_DIM];
let bounds = components
.bounds
.get(component as usize - 1)
.and_then(|o| o.as_ref())
.ok_or_else(|| "Missing bounds for given connected component")?;
for i in 0..(SPLIT_FACTOR * SPLIT_FACTOR) {
let (iy, ix) = i.div_rem_euclid(&SPLIT_FACTOR);
let x_start = bounds.x_min + ix * components.max_width / SPLIT_FACTOR;
let x_end = bounds.x_min + (ix + 1) * components.max_width / SPLIT_FACTOR;
let y_start = bounds.y_min + iy * components.max_height / SPLIT_FACTOR;
let y_end = bounds.y_min + (iy + 1) * components.max_height / SPLIT_FACTOR;
let mut count = 0;
for x in x_start..x_end {
for y in y_start..y_end {
if let Some(p) = components.components.get_pixel_checked(x, y)
&& p.0[0] == component
{
count += 1;
}
}
}
let size = (x_end + 1 - x_start) * (y_end + 1 - y_start);
if size == 0 {
return Err(format!(
"Got zero size for chunk [{x_start},{x_end}]x[{y_start},{y_end}]"
)
.into());
}
chunks[i as usize] = count as f32 / size as f32;
// print!("{} ", chunks[i as usize]);
// if i % SPLIT_FACTOR == SPLIT_FACTOR - 1 {
// print!("\n");
// }
}
let mut result = Self { chunks };
result.normalise();
Ok(result)
}
// }}}
// {{{ Distance
#[inline]
fn distance_squared_to(&self, other: &Self) -> f32 {
let mut total = 0.0;
for i in 0..IMAGE_VEC_DIM {
let d = self.chunks[i] - other.chunks[i];
total += d * d;
}
total
}
#[inline]
fn norm_squared(&self) -> f32 {
let mut total = 0.0;
for i in 0..IMAGE_VEC_DIM {
total += self.chunks[i] * self.chunks[i];
}
total
}
#[inline]
fn normalise(&mut self) {
let len = self.norm_squared().sqrt();
for i in 0..IMAGE_VEC_DIM {
self.chunks[i] /= len;
}
}
// }}}
}
// }}}
// {{{ Component bounds
#[derive(Clone, Copy)]
struct ComponentBounds {
x_min: u32,
y_min: u32,
x_max: u32,
y_max: u32,
}
struct ComponentsWithBounds {
components: ImageBuffer<Luma<u32>, Vec<u32>>,
// NOTE: the index is (the id of the component) - 1
// This is because the zero component represents the background,
// but we don't want to waste a place in this vector.
bounds: Vec<Option<ComponentBounds>>,
max_width: u32,
max_height: u32,
/// Stores the indices of `self.bounds` sorted based on their min position.
bounds_by_position: Vec<usize>,
}
impl ComponentsWithBounds {
fn from_image(image: &DynamicImage) -> Result<Self, Error> {
let image = threshold(&image.to_luma8(), 100, ThresholdType::Binary);
debug_image_buffer_log(&image)?;
let background = Luma([u8::MAX]);
let components = connected_components(&image, Connectivity::Eight, background);
let mut bounds: Vec<Option<ComponentBounds>> = Vec::new();
for x in 0..components.width() {
for y in 0..components.height() {
// {{{ Retrieve pixel if it's not backround
let component = components[(x, y)].0[0];
if component == 0 {
continue;
}
let index = component as usize - 1;
if index >= bounds.len() {
bounds.resize(index + 1, None);
}
// }}}
// {{{ Update bounds
if let Some(bounds) = (&mut bounds)[index].as_mut() {
bounds.x_min = bounds.x_min.min(x);
bounds.x_max = bounds.x_max.max(x);
bounds.y_min = bounds.y_min.min(y);
bounds.y_max = bounds.y_max.max(y);
} else {
bounds[index] = Some(ComponentBounds {
x_min: x,
x_max: x,
y_min: y,
y_max: y,
});
}
// }}}
}
}
// {{{ Remove components that are too large
for bound in &mut bounds {
if bound.map_or(false, |b| (b.x_max - b.x_min) >= 9 * image.width() / 10) {
*bound = None;
}
}
// }}}
// {{{ Compute max width/height
let max_width = bounds
.iter()
.filter_map(|o| o.as_ref())
.map(|b| b.x_max - b.x_min)
.max()
.ok_or_else(|| "No connected components found")?;
let max_height = bounds
.iter()
.filter_map(|o| o.as_ref())
.map(|b| b.y_max - b.y_min)
.max()
.ok_or_else(|| "No connected components found")?;
// }}}
let mut bounds_by_position: Vec<usize> = (0..(bounds.len()))
.filter(|i| bounds[*i].is_some())
.collect();
bounds_by_position.sort_by_key(|i| bounds[*i].unwrap().x_min);
Ok(Self {
components,
bounds,
max_width,
max_height,
bounds_by_position,
})
}
}
// }}}
// {{{ Char measurements
pub struct CharMeasurements {
chars: Vec<(char, ComponentVec)>,
}
impl CharMeasurements {
// {{{ Creation
pub fn from_text(face: &mut Face, string: &str, _weight: Option<u32>) -> Result<Self, Error> {
// These are bad estimates lol
let char_w = 35;
let char_h = 60;
let mut canvas = BitmapCanvas::new(10 + char_w * string.len() as u32, char_h + 10);
canvas.text(
(5, 5),
face,
TextStyle {
stroke: None,
drop_shadow: None,
align: (Align::Start, Align::Start),
size: char_h,
color: Color::BLACK,
// TODO: do we want to use the weight hint for resilience?
weight: None,
},
&string,
)?;
let buffer = ImageBuffer::from_raw(canvas.width, canvas.height(), canvas.buffer.to_vec())
.ok_or_else(|| "Failed to turn buffer into canvas")?;
let image = DynamicImage::ImageRgb8(buffer);
debug_image_log(&image)?;
let components = ComponentsWithBounds::from_image(&image)?;
let mut chars = Vec::with_capacity(string.len());
for (i, char) in string.chars().enumerate() {
chars.push((
char,
ComponentVec::from_component(
&components,
components.bounds_by_position[i] as u32 + 1,
)?,
))
}
Ok(Self { chars })
}
// }}}
// {{{ Recognition
pub fn recognise(&self, image: &DynamicImage) -> Result<String, Error> {
let components = ComponentsWithBounds::from_image(image)?;
let mut result = String::new();
for i in &components.bounds_by_position {
let vec = ComponentVec::from_component(&components, *i as u32 + 1)?;
let best_match = self
.chars
.iter()
.map(|(i, v)| (*i, v, v.distance_squared_to(&vec)))
.min_by(|(_, _, d1), (_, _, d2)| {
d1.partial_cmp(d2).expect("NaN distance encountered")
})
.map(|(i, _, d)| (d.sqrt(), i))
.ok_or_else(|| "No chars in cache")?;
// println!("char '{}', distance {}", best_match.1, best_match.0);
if best_match.0 <= (IMAGE_VEC_DIM * 10) as f32 {
result.push(best_match.1);
}
}
Ok(result)
}
// }}}
}
// }}}