working GUI initial
This commit is contained in:
commit
aa1c01a6dd
|
@ -0,0 +1 @@
|
|||
/target
|
File diff suppressed because it is too large
Load Diff
|
@ -0,0 +1,22 @@
|
|||
[package]
|
||||
name = "supper"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
|
||||
eframe = "0.18"
|
||||
egui_extras = { version = "0.18", features = ["image"] }
|
||||
ehttp = "0.2"
|
||||
image = { version = "0.24" }
|
||||
poll-promise = "0.1"
|
||||
scrap = "0.5"
|
||||
|
||||
anyhow = "1.0"
|
||||
futures = "0.3"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
reqwest = { version = "0.11", features = ["json"] }
|
|
@ -0,0 +1,42 @@
|
|||
use std::{time::Duration, thread};
|
||||
|
||||
use anyhow::Result;
|
||||
use image::{RgbImage, Rgb};
|
||||
use scrap::{Capturer, Display};
|
||||
|
||||
fn get_raw_frame(capturer: &mut Capturer) -> Result<Vec<u8>> {
|
||||
loop {
|
||||
let frame = capturer.frame();
|
||||
match frame {
|
||||
Ok(buffer) => return Ok(buffer.to_vec()),
|
||||
Err(err) => {
|
||||
if err.kind() == std::io::ErrorKind::WouldBlock {
|
||||
thread::sleep(Duration::from_secs(1) / 60);
|
||||
continue;
|
||||
} else {
|
||||
anyhow::bail!("Failed to capture frame")
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_frame() -> Result<RgbImage> {
|
||||
let mut capturer = Capturer::new(Display::primary()?)?;
|
||||
let frame = get_raw_frame(&mut capturer)?;
|
||||
let mut image = RgbImage::new(capturer.width() as u32, capturer.height() as u32);
|
||||
|
||||
let stride = frame.len() / capturer.height();
|
||||
for y in 0..capturer.height() {
|
||||
for x in 0..capturer.width() {
|
||||
let i = stride * y + 4 * x;
|
||||
image.put_pixel(
|
||||
x as u32,
|
||||
y as u32,
|
||||
Rgb([frame[i + 2], frame[i + 1], frame[i]]),
|
||||
);
|
||||
}
|
||||
|
||||
}
|
||||
Ok(image)
|
||||
}
|
|
@ -0,0 +1,59 @@
|
|||
use std::{collections::HashMap, time::Duration};
|
||||
|
||||
use anyhow::Result;
|
||||
use egui_extras::RetainedImage;
|
||||
use scrap::{Capturer, Display};
|
||||
|
||||
use crate::{
|
||||
capture,
|
||||
image_processing::{self, Region},
|
||||
ocr,
|
||||
state::{ParsedFrame, SharedAppState},
|
||||
};
|
||||
|
||||
async fn run_loop_once(state: &SharedAppState, regions: &[Region]) -> Result<()> {
|
||||
let frame = capture::get_frame()?;
|
||||
let ocr_results = ocr::ocr_all_regions(&frame, ®ions).await;
|
||||
|
||||
let mut saved_frames = HashMap::new();
|
||||
if state.lock().unwrap().debug_frames {
|
||||
for region in regions {
|
||||
let mut extracted = image_processing::extract_region(&frame, region);
|
||||
image_processing::filter_to_white(&mut extracted, 0.95, 0.05);
|
||||
let retained = RetainedImage::from_image_bytes(
|
||||
®ion.name,
|
||||
&image_processing::to_png_bytes(&extracted),
|
||||
)
|
||||
.unwrap();
|
||||
saved_frames.insert(region.name.clone(), retained);
|
||||
}
|
||||
}
|
||||
{
|
||||
let mut state = state.lock().unwrap();
|
||||
let parsed = ParsedFrame::parse(&ocr_results);
|
||||
if parsed.lap_time.is_some() {
|
||||
state.last_frame = Some(parsed);
|
||||
}
|
||||
state.raw_data = ocr_results;
|
||||
state.saved_frames = saved_frames;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn run_control_loop(state: SharedAppState) -> Result<()> {
|
||||
let regions = vec![
|
||||
Region::parse("health 91 1364 52 24").unwrap(),
|
||||
Region::parse("gas 208 1364 52 24").unwrap(),
|
||||
Region::parse("tyres 325 1364 52 24").unwrap(),
|
||||
Region::parse("lap 2295 46 140 87").unwrap(),
|
||||
Region::parse("best 2325 169 183 43").unwrap(),
|
||||
Region::parse("lap_time 2325 222 183 43").unwrap(),
|
||||
];
|
||||
|
||||
loop {
|
||||
if let Err(e) = run_loop_once(&state, ®ions).await {
|
||||
eprintln!("Error in control loop: {:?}", e)
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
|
@ -0,0 +1,75 @@
|
|||
use anyhow::Result;
|
||||
use image::{RgbImage, Rgb, codecs::png::PngEncoder, ColorType, ImageEncoder};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Region {
|
||||
pub name: String,
|
||||
x: usize,
|
||||
y: usize,
|
||||
width: usize,
|
||||
height: usize
|
||||
}
|
||||
|
||||
impl Region {
|
||||
pub fn parse(encoded: &str) -> Result<Region> {
|
||||
let split: Vec<_> = encoded.split(' ').collect();
|
||||
Ok(Region {
|
||||
name: split.get(0).ok_or(anyhow::anyhow!("expected name in region"))?.to_string(),
|
||||
x: split.get(1).ok_or(anyhow::anyhow!("failed to parse x from region"))?.parse()?,
|
||||
y: split.get(2).ok_or(anyhow::anyhow!("failed to parse y from region"))?.parse()?,
|
||||
width: split.get(3).ok_or(anyhow::anyhow!("failed to parse width from region"))?.parse()?,
|
||||
height: split.get(4).ok_or(anyhow::anyhow!("failed to parse height from region"))?.parse()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_region(
|
||||
image: &RgbImage,
|
||||
region: &Region,
|
||||
) -> RgbImage {
|
||||
let mut buffer = RgbImage::new(region.width as u32, region.height as u32);
|
||||
for y in 0..region.height {
|
||||
for x in 0..region.width {
|
||||
buffer.put_pixel(
|
||||
x as u32,
|
||||
y as u32,
|
||||
*image.get_pixel((region.x + x) as u32, (region.y + y) as u32)
|
||||
);
|
||||
}
|
||||
}
|
||||
buffer
|
||||
}
|
||||
|
||||
pub fn filter_to_white(
|
||||
image: &mut RgbImage,
|
||||
threshold: f64,
|
||||
variance_threshold: f64
|
||||
) {
|
||||
let past_threshold_color = |v: u8| v as f64 >= (u8::MAX as f64 * threshold);
|
||||
let color_diff = |a: u8, b: u8| (a.abs_diff(b) as f64) / (u8::MAX as f64);
|
||||
for y in 0..image.height() {
|
||||
for x in 0..image.width() {
|
||||
let pixel = image.get_pixel(x, y);
|
||||
let [r, g, b] = pixel.0;
|
||||
if past_threshold_color(r) && past_threshold_color(g) && past_threshold_color(b) && color_diff(r, g) < variance_threshold && color_diff(r, b) < variance_threshold && color_diff(b, g) < variance_threshold {
|
||||
// This is a white pixel, make it black
|
||||
image.put_pixel(x, y, Rgb([0, 0, 0]));
|
||||
} else {
|
||||
// Make it white
|
||||
image.put_pixel(x, y, Rgb([255, 255, 255]))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_png_bytes(image: &RgbImage) -> Vec<u8> {
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
let encoder = PngEncoder::new(&mut buffer);
|
||||
encoder.write_image(
|
||||
image.as_raw(),
|
||||
image.width(),
|
||||
image.height(),
|
||||
ColorType::Rgb8
|
||||
).expect("failed encoding image to PNG");
|
||||
buffer
|
||||
}
|
|
@ -0,0 +1,102 @@
|
|||
// #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] // hide console window on Windows in release
|
||||
|
||||
mod capture;
|
||||
mod control_loop;
|
||||
mod image_processing;
|
||||
mod ocr;
|
||||
mod state;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use eframe::egui;
|
||||
use egui_extras::RetainedImage;
|
||||
use state::{AppState, SharedAppState};
|
||||
|
||||
#[tokio::main(flavor = "multi_thread", worker_threads = 8)]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let state = Arc::new(Mutex::new(AppState::default()));
|
||||
{
|
||||
let state = state.clone();
|
||||
let _ = tokio::spawn(async move {
|
||||
control_loop::run_control_loop(state)
|
||||
.await
|
||||
.expect("control loop failed");
|
||||
});
|
||||
}
|
||||
|
||||
let options = eframe::NativeOptions::default();
|
||||
eframe::run_native(
|
||||
"Supper OCR",
|
||||
options,
|
||||
Box::new(|_cc| Box::new(MyApp::new(state))),
|
||||
);
|
||||
}
|
||||
|
||||
struct MyApp {
|
||||
state: SharedAppState,
|
||||
}
|
||||
|
||||
impl MyApp {
|
||||
pub fn new(state: SharedAppState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
impl eframe::App for MyApp {
|
||||
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
egui::SidePanel::left("frame").show(ctx, |ui| {
|
||||
if let Some(frame) = &state.last_frame {
|
||||
ui.heading("Race data");
|
||||
ui.label(format!("Lap: {}", frame.lap.unwrap_or(0)));
|
||||
ui.separator();
|
||||
ui.label(format!("Health: {}/100", frame.health.unwrap_or(0)));
|
||||
ui.label(format!("Gas: {}/100", frame.gas.unwrap_or(0)));
|
||||
ui.label(format!("Tyres: {}/100", frame.tyres.unwrap_or(0)));
|
||||
ui.separator();
|
||||
ui.label(format!(
|
||||
"Lap time: {}s",
|
||||
frame
|
||||
.lap_time
|
||||
.unwrap_or(Duration::from_secs(0))
|
||||
.as_secs_f32()
|
||||
));
|
||||
ui.label(format!(
|
||||
"Best lap: {}s",
|
||||
frame
|
||||
.best_time
|
||||
.unwrap_or(Duration::from_secs(0))
|
||||
.as_secs_f32()
|
||||
));
|
||||
}
|
||||
|
||||
ui.separator();
|
||||
ui.heading("Raw OCR results");
|
||||
let mut raw_data_sorted: Vec<_> = state.raw_data.iter().collect();
|
||||
raw_data_sorted.sort();
|
||||
for (key, val) in raw_data_sorted {
|
||||
ui.label(format!("{}: {:?}", key, val));
|
||||
}
|
||||
|
||||
ui.separator();
|
||||
ui.checkbox(&mut state.debug_frames, "Debug OCR regions");
|
||||
});
|
||||
egui::CentralPanel::default().show(ctx, |ui| {});
|
||||
|
||||
if state.debug_frames {
|
||||
egui::SidePanel::right("screenshots").show(ctx, |ui| {
|
||||
let mut screenshots_sorted: Vec<_> = state.saved_frames.iter().collect();
|
||||
screenshots_sorted.sort_by_key(|(name, _)| name.clone());
|
||||
for (name, image) in screenshots_sorted {
|
||||
ui.label(name);
|
||||
image.show_max_size(ui, ui.available_size());
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
ctx.request_repaint();
|
||||
}
|
||||
}
|
|
@ -0,0 +1,71 @@
|
|||
use std::{collections::HashMap, sync::{Arc, Mutex}};
|
||||
|
||||
use anyhow::Result;
|
||||
use image::RgbImage;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::image_processing::{Region, extract_region, filter_to_white};
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug)]
|
||||
pub struct OcrRegion {
|
||||
pub confidence: f64,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug)]
|
||||
struct OcrResult {
|
||||
regions: Vec<OcrRegion>,
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
async fn run_ocr(image: &RgbImage) -> Result<Vec<OcrRegion>> {
|
||||
let client = reqwest::Client::new();
|
||||
let response = client
|
||||
.post("http://localhost:3000/")
|
||||
.body(crate::image_processing::to_png_bytes(&image))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
anyhow::bail!("failed to run OCR query")
|
||||
}
|
||||
let result: OcrResult = response.json().await?;
|
||||
Ok(result.regions)
|
||||
}
|
||||
|
||||
pub async fn ocr_all_regions(image: &RgbImage, regions: &[Region]) -> HashMap<String, Option<String>> {
|
||||
let results = Arc::new(Mutex::new(HashMap::new()));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for region in regions {
|
||||
let filtered_image = extract_region(image, region);
|
||||
let region = region.clone();
|
||||
let results = results.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
let mut image = filtered_image;
|
||||
filter_to_white(&mut image, 0.95, 0.05);
|
||||
let ocr_results = run_ocr(&image).await;
|
||||
let value = match ocr_results {
|
||||
Ok(ocr_regions) => {
|
||||
if ocr_regions.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let mut out = String::new();
|
||||
for r in ocr_regions {
|
||||
out += &r.value;
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
}
|
||||
Err(_) => None
|
||||
};
|
||||
results.lock().unwrap().insert(region.name, value);
|
||||
}));
|
||||
}
|
||||
for handle in handles {
|
||||
handle.await.expect("failed to join task in OCR");
|
||||
}
|
||||
|
||||
let results = results.lock().unwrap().clone();
|
||||
results
|
||||
}
|
|
@ -0,0 +1,66 @@
|
|||
use std::{sync::{Arc, Mutex}, time::Duration, collections::HashMap};
|
||||
|
||||
use egui_extras::RetainedImage;
|
||||
|
||||
|
||||
pub struct ParsedFrame {
|
||||
pub lap: Option<usize>,
|
||||
|
||||
pub health: Option<usize>,
|
||||
pub gas: Option<usize>,
|
||||
pub tyres: Option<usize>,
|
||||
|
||||
pub best_time: Option<Duration>,
|
||||
pub lap_time: Option<Duration>,
|
||||
}
|
||||
|
||||
fn parse_duration(time: &str) -> Option<Duration> {
|
||||
let sep = time.find(':')?;
|
||||
let (minutes, secs) = time.split_at(sep);
|
||||
|
||||
println!("{} {}", minutes, secs);
|
||||
let minutes = minutes.parse::<f64>().ok()?;
|
||||
let secs = secs[1..secs.len()].parse::<f64>().ok()?;
|
||||
|
||||
Some(Duration::from_secs_f64(60.0 * minutes + secs))
|
||||
}
|
||||
|
||||
fn parse_to_duration(time: Option<&Option<String>>) -> Option<Duration> {
|
||||
parse_duration(&time?.clone()?)
|
||||
}
|
||||
|
||||
fn check_0_100(v: usize) -> Option<usize> {
|
||||
if v > 100 {
|
||||
None
|
||||
} else {
|
||||
Some(v)
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_to_0_100(v: Option<&Option<String>>) -> Option<usize> {
|
||||
check_0_100(v?.clone()?.parse::<usize>().ok()?)
|
||||
}
|
||||
|
||||
impl ParsedFrame {
|
||||
pub fn parse(raw: &HashMap<String, Option<String>>) -> Self {
|
||||
Self {
|
||||
lap: parse_to_0_100(raw.get("lap")),
|
||||
health: parse_to_0_100(raw.get("health")),
|
||||
gas: parse_to_0_100(raw.get("gas")),
|
||||
tyres: parse_to_0_100(raw.get("tyres")),
|
||||
best_time: parse_to_duration(raw.get("best")),
|
||||
lap_time: parse_to_duration(raw.get("lap_time")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct AppState {
|
||||
pub raw_data: HashMap<String, Option<String>>,
|
||||
pub last_frame: Option<ParsedFrame>,
|
||||
|
||||
pub debug_frames: bool,
|
||||
pub saved_frames: HashMap<String, RetainedImage>,
|
||||
}
|
||||
|
||||
pub type SharedAppState = Arc<Mutex<AppState>>;
|
Loading…
Reference in New Issue