Simmer/Waywrite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Connect ocrs to new ui

Browse Source
This commit is contained in:
Ethan Simmons
2024-05-05 18:02:54 -05:00
parent 092c44e61e
commit 392cea7a6a
9 changed files with 306 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
src/process_point.rs
View File

@@ -5,51 +5,19 @@ use std::time::Instant;
use rten_tensor::{NdTensor, AsView};
use rten::Model;
use ocrs::{OcrEngine, OcrEngineParams};
use image::{ColorType, ImageFormat};
const MATRIX_LEN: usize = 800;
const TEXT_MATRIX_RATIO: f64 = 0.5;
const LINE_WIDTH: f64 = 10.0;
const MATRIX_X_SIZE: f64 = MATRIX_LEN as f64;
const TEXT_X_SIZE: f64 = (MATRIX_X_SIZE as f64) * TEXT_MATRIX_RATIO;
const TEXT_X_OFFSET: f64 = (MATRIX_X_SIZE - TEXT_X_SIZE) / 2.0;
const LINE_WIDTH_X_OFFSET: f64 = LINE_WIDTH / 2.0;
#[derive(Debug)]
pub struct Point {
pub x: f64,
pub y: f64,
pub new_line: bool,
}
pub fn print_words(points: &Vec<Point>) -> Result<(), Box<dyn Error>> {
pub fn print_words(image: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
let begin = Instant::now();
let processed_data = process(points);
println!("{:#?}", begin.elapsed());
let begin = Instant::now();
ocr(processed_data)?;
ocr(image).unwrap();
println!("{:#?}", begin.elapsed());
Ok(())
}
fn process(points: &Vec<Point>) -> NdTensor<f32, 3> {
let matrix = to_matrix(points);
let y_len = matrix[0].len();
let x_len = matrix[0][0].len();
let image_data: Box<[u8]> = matrix.iter().flatten().flatten().map(|f| if *f > 0.5 { u8::from(0) } else { u8::from(255) }).collect();
let data: Vec<f32> = matrix.into_iter().flatten().flatten().map(|f| f as f32).collect();
image::save_buffer_with_format("./image.png", &image_data, x_len as u32, y_len as u32, ColorType::L8, ImageFormat::Png).unwrap();
NdTensor::from_data([1, y_len, x_len], data)
}
fn ocr(data: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
let detection_model_data = fs::read("text-detection.rten")?;
@@ -58,7 +26,6 @@ fn ocr(data: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
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),
@@ -87,91 +54,3 @@ fn ocr(data: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
}
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<Point>) -> Vec<Vec<Vec<f64>>> {
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 matrix_y_size = MATRIX_X_SIZE * y_ratio;
let text_y_size = TEXT_X_SIZE * y_ratio;
let text_y_offset = TEXT_X_OFFSET * y_ratio;
let line_width_y_offset = LINE_WIDTH_X_OFFSET * y_ratio;
let x_scale = MATRIX_X_SIZE / x_len;
let y_scale = matrix_y_size / y_len;
let mut matrix: Vec<Vec<f64>> = vec![
vec![0.0; MATRIX_LEN]; (matrix_y_size as usize) + 1
];
let scaled_points: Vec<((f64, f64), bool)> = points
.iter()
.map(|point| {
let x_scaled = ((point.x - min_x) * x_scale) + TEXT_X_OFFSET;
let y_scaled = ((point.y - min_y) * y_scale) + text_y_offset;
((x_scaled, y_scaled), point.new_line)
}).collect();
let mut last_x = 0.0;
let mut last_y = 0.0;
for ((current_x, current_y), newline) in scaled_points {
if !newline {
let curr_x_start = current_x - LINE_WIDTH_X_OFFSET;
let curr_x_end = current_x + LINE_WIDTH_X_OFFSET;
let last_x_start = last_x - LINE_WIDTH_X_OFFSET;
let last_x_end = last_x + LINE_WIDTH_X_OFFSET;
let top_y = current_y.max(last_y) + line_width_y_offset;
let bottom_y = current_y.min(last_y) - line_width_y_offset;
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, current_y));
let right_line_y = line(x as f64, (last_x_end, last_y), (curr_x_end, current_x));
let top_line_y = left_line_y
.max(right_line_y)
.min(top_y) as usize;
let bottom_line_y = left_line_y
.min(right_line_y)
.max(bottom_y) as usize;
for y in bottom_line_y..(top_line_y + 1) {
matrix[y][x] = 1.0;
}
}
}
last_x = current_x;
last_y = current_y;
}
vec![matrix]
}

4
src/wayland.rs
View File

@@ -1,6 +1,6 @@
mod ui;
pub mod app;
mod draw;
pub mod draw;
mod ui;
mod builder;
#[derive(Clone, Copy, Debug)]

74
src/wayland/app.rs
View File

@@ -32,17 +32,21 @@ use wayland_client::{
use std::time::Duration;
use rten_tensor::NdTensor;
use anyhow::Result;
use crate::wayland::{
draw::DrawPath,
draw::{DrawPath, Draw},
builder::{TextAreaBuilder, DrawAreaBuilder, ButtonBuilder},
ui::{self, Widget},
ui::{self, Widget, ButtonType},
Position,
};
const WINDOW_HEIGHT: u32 = 256;
const WINDOW_WIDTH: u32 = 512;
use crate::process_point::print_words;
const WINDOW_HEIGHT: usize = 256;
const WINDOW_WIDTH: usize = 512;
pub(crate) struct SimpleWindow {
registry_state: RegistryState,
@@ -55,8 +59,8 @@ pub(crate) struct SimpleWindow {
pool: SlotPool,
window: Window,
ui: Option<ui::Window>,
width: u32,
height: u32,
width: usize,
height: usize,
cursor_down: bool,
exit: bool,
first_configure: bool,
@@ -87,8 +91,8 @@ impl SimpleWindow {
window.set_title("A window");
window.set_app_id("simmer.simplewindow");
window.set_min_size(Some((WINDOW_WIDTH, WINDOW_HEIGHT)));
window.set_max_size(Some((WINDOW_WIDTH, WINDOW_HEIGHT)));
window.set_min_size(Some((WINDOW_WIDTH as u32, WINDOW_HEIGHT as u32)));
window.set_max_size(Some((WINDOW_WIDTH as u32, WINDOW_HEIGHT as u32)));
window.commit();
@@ -260,8 +264,8 @@ impl WindowHandler for SimpleWindow {
let mut draw_area = DrawAreaBuilder::new();
draw_area.position(0.0, 50.0);
draw_area.width(self.width as f32 - 100.0);
draw_area.height(self.height as f32 - 50.0);
draw_area.width(self.width - 100);
draw_area.height(self.height - 50);
draw_area.path(&mut self.draw_path);
let draw_area = draw_area.finish();
@@ -273,22 +277,31 @@ impl WindowHandler for SimpleWindow {
draw_area,
);
let mut button = ButtonBuilder::new();
let mut enter_button = ButtonBuilder::new();
button.position(self.width as f32 - 100.0, 50.0);
button.width(100.0);
button.height(50.0);
let button = button.finish();
enter_button.button_type(ButtonType::Enter);
enter_button.position(self.width - 100, 60);
enter_button.width(100);
enter_button.height(50);
let enter_button = enter_button.finish();
let mut clear_button = ButtonBuilder::new();
clear_button.button_type(ButtonType::Clear);
clear_button.position(self.width - 100, 120);
clear_button.width(100);
clear_button.height(50);
let clear_button = clear_button.finish();
let mut text_area = TextAreaBuilder::new();
text_area.position(0.0, 0.0);
text_area.width(self.width as f32);
text_area.height(50.0);
text_area.width(self.width);
text_area.height(50);
let text_area = text_area.finish();
window.add_widget(button);
window.add_widget(enter_button);
window.add_widget(clear_button);
window.add_widget(text_area);
self.ui = Some(window);
@@ -354,12 +367,25 @@ impl PointerHandler for SimpleWindow {
window.draw_area.path.add_point(pos, true);
}
for widget in &window.widgets {
if let Widget::Button(button) = widget {
if button.contains_point(pos) {
window.draw_area.path = DrawPath::default();
}
}
for button in window.widgets
.iter()
.filter_map(|widget| {
if let Widget::Button(button) = widget {
Some(button)
} else {
None
}
})
{
match button.button_type {
ButtonType::Clear if button.contains_point(pos) => {
window.draw_area.path = DrawPath::default();
},
ButtonType::Enter if button.contains_point(pos) => {
print_words((&window.draw_area).into()).unwrap()
},
_ => (),
}
}
}
}

49
src/wayland/builder.rs
View File

@@ -1,5 +1,5 @@
use crate::wayland::{
ui::{TextArea, TextAreaColors, Button, ButtonColors, Widget},
ui::{TextArea, TextAreaColors, Button, ButtonColors, Widget, ButtonType},
draw::{DrawAreaColors, DrawArea, DrawPath},
Position
};
@@ -7,8 +7,8 @@ use crate::wayland::{
use raqote::{DrawOptions, PathBuilder};
pub(super) struct TextAreaBuilder {
pub(super) width: Option<f32>,
pub(super) height: Option<f32>,
pub(super) width: Option<usize>,
pub(super) height: Option<usize>,
pub(super) position: Option<Position>,
pub(super) colors: Option<TextAreaColors>,
pub(super) options: Option<DrawOptions>,
@@ -25,11 +25,11 @@ impl<'a> TextAreaBuilder {
}
}
pub(super) fn width(&mut self, width: f32) {
pub(super) fn width(&mut self, width: usize) {
self.width = Some(width)
}
pub(super) fn height(&mut self, height: f32) {
pub(super) fn height(&mut self, height: usize) {
self.height = Some(height)
}
@@ -56,11 +56,12 @@ impl<'a> TextAreaBuilder {
}
pub(super) struct ButtonBuilder {
pub(super) width: Option<f32>,
pub(super) height: Option<f32>,
pub(super) width: Option<usize>,
pub(super) height: Option<usize>,
pub(super) position: Option<Position>,
pub(super) colors: Option<ButtonColors>,
pub(super) options: Option<DrawOptions>,
pub(super) button_type: Option<ButtonType>,
}
impl ButtonBuilder {
@@ -71,25 +72,30 @@ impl ButtonBuilder {
position: None,
colors: None,
options: None,
button_type: None,
}
}
pub(super) fn width(&mut self, width: f32) {
pub(super) fn width(&mut self, width: usize) {
self.width = Some(width)
}
pub(super) fn height(&mut self, height: f32) {
pub(super) fn height(&mut self, height: usize) {
self.height = Some(height)
}
pub(super) fn position(&mut self, x: f32, y: f32) {
self.position = Some(Position::new(x, y))
pub(super) fn position(&mut self, x: usize, y: usize) {
self.position = Some(Position::new(x as f32, y as f32))
}
pub(super) fn colors(&mut self, colors: ButtonColors) {
self.colors = Some(colors)
}
pub(super) fn button_type(&mut self, button_type: ButtonType) {
self.button_type = Some(button_type)
}
pub(super) fn finish(self) -> Widget {
let width = self.width.expect("Button must have a width");
let height = self.height.expect("Button must have a height");
@@ -97,14 +103,15 @@ impl ButtonBuilder {
let mut button_path = PathBuilder::new();
button_path.move_to(position.x, position.y);
button_path.line_to(position.x + width, position.y);
button_path.line_to(position.x + width, position.y + height);
button_path.line_to(position.x, position.y + height);
button_path.line_to(position.x + width as f32, position.y);
button_path.line_to(position.x + width as f32, position.y + height as f32);
button_path.line_to(position.x, position.y + height as f32);
button_path.line_to(position.x, position.y);
let button_path = button_path.finish();
Widget::Button(
Button {
button_type: self.button_type.expect("Button must have a type"),
width: self.width.expect("Button must have width"),
height: self.height.expect("Button must have height"),
position: self.position.expect("Button must have a position"),
@@ -117,8 +124,8 @@ impl ButtonBuilder {
}
pub(super) struct DrawAreaBuilder {
pub(super) width: Option<f32>,
pub(super) height: Option<f32>,
pub(super) width: Option<usize>,
pub(super) height: Option<usize>,
pub(super) position: Option<Position>,
pub(super) path: Option<DrawPath>,
pub(super) colors: Option<DrawAreaColors>,
@@ -137,11 +144,11 @@ impl DrawAreaBuilder {
}
}
pub(super) fn width(&mut self, width: f32) {
pub(super) fn width(&mut self, width: usize) {
self.width = Some(width)
}
pub(super) fn height(&mut self, height: f32) {
pub(super) fn height(&mut self, height: usize) {
self.height = Some(height)
}
@@ -168,9 +175,9 @@ impl DrawAreaBuilder {
let mut draw_rect_path = PathBuilder::new();
draw_rect_path.move_to(position.x, position.y);
draw_rect_path.line_to(position.x + width, position.y);
draw_rect_path.line_to(position.x + width, position.y + height);
draw_rect_path.line_to(position.x, position.y + height);
draw_rect_path.line_to(position.x + width as f32, position.y);
draw_rect_path.line_to(position.x + width as f32, position.y + height as f32);
draw_rect_path.line_to(position.x, position.y + height as f32);
draw_rect_path.line_to(position.x, position.y);
let draw_rect_path = draw_rect_path.finish();

96
src/wayland/draw.rs
View File

@@ -1,6 +1,8 @@
use raqote::{DrawTarget, Color, StrokeStyle, PathBuilder, Source, DrawOptions};
use raqote::{DrawTarget, Color, StrokeStyle, PathBuilder, Source, DrawOptions, SolidSource};
use rten_tensor::NdTensor;
use crate::wayland::Position;
use crate::process_point::print_words;
pub(super) trait Draw {
fn render(&self, area: &mut DrawTarget<&mut [u32]>);
@@ -9,20 +11,20 @@ pub(super) trait Draw {
type Path = Vec<Position>;
#[derive(Default)]
pub(super) struct DrawPath {
pub(super) paths: Vec<Path>,
pub(super) newline: bool,
pub(crate) struct DrawPath {
pub(crate) paths: Vec<Path>,
pub(crate) newline: bool,
}
impl DrawPath {
pub(super) fn new() -> Self {
pub(crate) fn new() -> Self {
Self {
paths: Vec::new(),
newline: false,
}
}
pub(super) fn add_point(&mut self, point: Position, newline: bool) {
pub(crate) fn add_point(&mut self, point: Position, newline: bool) {
if newline {
self.paths.push(vec![point])
} else if self.paths.len() > 0 {
@@ -33,15 +35,15 @@ impl DrawPath {
}
}
pub(super) fn into_iter(self) -> std::vec::IntoIter<Path> {
pub(crate) fn into_iter(self) -> std::vec::IntoIter<Path> {
self.paths.into_iter()
}
pub(super) fn iter(&self) -> std::slice::Iter<Path> {
pub(crate) fn iter(&self) -> std::slice::Iter<Path> {
self.paths.iter()
}
pub(super) fn iter_mut(&mut self) -> std::slice::IterMut<Path> {
pub(crate) fn iter_mut(&mut self) -> std::slice::IterMut<Path> {
self.paths.iter_mut()
}
}
@@ -63,8 +65,8 @@ impl Default for DrawAreaColors {
}
pub(super) struct DrawArea {
pub(super) width: f32,
pub(super) height: f32,
pub(super) width: usize,
pub(super) height: usize,
pub(super) position: Position,
pub(super) path: DrawPath,
pub(super) shape: raqote::Path,
@@ -89,6 +91,78 @@ impl DrawArea {
self.shape.contains_point(0.1, point.x, point.y)
}
}
impl From<&DrawArea> for NdTensor<f32, 3> {
fn from(draw_area: &DrawArea) -> Self {
let mut target = DrawTarget::new(draw_area.width as i32, draw_area.height as i32);
target.fill_rect(
0.0,
0.0,
draw_area.width as f32,
draw_area.height as f32,
&Source::Solid(SolidSource::from_unpremultiplied_argb(0xff, 0xff, 0xff, 0xff)),
&DrawOptions::new(),
);
let mut draw_path = PathBuilder::new();
for path in draw_area.path.iter() {
if path.len() > 1 {
let mut path = path
.into_iter()
.map(|position| Position::new(position.x - draw_area.position.x, position.y - draw_area.position.y));
let first_point = path
.by_ref()
.next();
if let Some(point) = first_point {
draw_path.move_to(point.x, point.y);
} else {
break;
}
for point in path {
draw_path.line_to(point.x, point.y);
}
}
}
let draw_path = draw_path.finish();
target.stroke(
&draw_path,
&Source::Solid(SolidSource::from_unpremultiplied_argb(0xff, 0x00, 0x00, 0x00)),
&StrokeStyle {width: 3.0, ..StrokeStyle::default()},
&DrawOptions::new(),
);
let data: Vec<f32> = target
.get_data_u8()
.chunks(4)
.into_iter()
.map(|chunk| chunk[0..3].into_iter().map(|byte| *byte as f32).sum::<f32>() / 765.0)
.collect();
image::save_buffer_with_format(
"./image.png",
&data
.iter()
.map(|num| (num * 255.0).floor() as u8)
.collect::<Vec<u8>>(),
draw_area.width as u32,
draw_area.height as u32,
image::ColorType::L8,
image::ImageFormat::Png
).unwrap();
NdTensor::from_data([1, draw_area.width as usize, draw_area.height as usize], data)
}
}
impl Draw for DrawArea {

19
src/wayland/ui.rs
View File

@@ -103,9 +103,16 @@ impl Default for ButtonColors {
}
}
#[derive(PartialEq, Eq)]
pub(super) enum ButtonType {
Clear,
Enter,
}
pub(super) struct Button {
pub(super) width: f32,
pub(super) height: f32,
pub(super) button_type: ButtonType,
pub(super) width: usize,
pub(super) height: usize,
pub(super) position: Position,
pub(super) shape: raqote::Path,
pub(super) colors: ButtonColors,
@@ -149,8 +156,8 @@ impl Default for TextAreaColors {
}
pub(super) struct TextArea {
pub(super) width: f32,
pub(super) height: f32,
pub(super) width: usize,
pub(super) height: usize,
pub(super) position: Position,
pub(super) colors: TextAreaColors,
pub(super) options: DrawOptions,
@@ -172,8 +179,8 @@ impl Draw for TextArea {
area.fill_rect(
self.position.x,
self.position.y,
self.width,
self.height,
self.width as f32,
self.height as f32,
&self.background_source(),
&self.options,
);