diff --git a/config.json b/config.json index 3b82065..56f9e4e 100644 --- a/config.json +++ b/config.json @@ -46,18 +46,22 @@ "use_ocr_cache": null }, { - "name": "position", - "x": 50, - "y": 44, - "width": 250, - "height": 100 + "name": "position", + "x": 50, + "y": 44, + "width": 250, + "height": 100, + "threshold": null, + "use_ocr_cache": null }, { - "name": "lap", - "x": 2290, - "y": 44, - "width": 230, - "height": 100 + "name": "lap", + "x": 2290, + "y": 44, + "width": 230, + "height": 100, + "threshold": null, + "use_ocr_cache": null } ], "track_region": { @@ -69,6 +73,20 @@ "threshold": 0.85, "use_ocr_cache": null }, + "penalty_orange_region": { + "name": "penalty", + "x": 993, + "y": 120, + "width": 50, + "height": 30, + "threshold": 0.9, + "use_ocr_cache": null + }, + "penalty_orange_color": [ + 255, + 177, + 0 + ], "ocr_interval_ms": 500, "track_recognition_threshold": 10, "dump_frame_fraction": null, diff --git a/src/analysis.rs b/src/analysis.rs index f15d18e..13d0742 100644 --- a/src/analysis.rs +++ b/src/analysis.rs @@ -6,7 +6,7 @@ use std::{ use anyhow::Result; use egui_extras::RetainedImage; -use image::RgbImage; +use image::{Rgb, RgbImage}; use scrap::{Capturer, Display}; @@ -19,6 +19,24 @@ use crate::{ state::{AppState, DebugOcrFrame, LapState, RaceState, SharedAppState}, }; +fn check_penalty(image: &RgbImage, config: &Config) -> bool { + let region = match config.penalty_orange_region.as_ref() { + Some(s) => s, + None => return false, + }; + let color = match config.penalty_orange_color.as_ref() { + Some(c) => c, + None => return false, + }; + + let fraction = image_processing::check_target_color_fraction( + image, + region, + Rgb([color.0, color.1, color.2]), + ); + fraction >= region.threshold.unwrap_or(0.90) +} + fn is_finished_lap(state: &AppState, frame: &LapState) -> bool { if let Some(race) = &state.current_race { if let Some(last_finish) = &race.last_lap_record_time { @@ -72,6 +90,8 @@ fn handle_new_frame(state: &mut AppState, lap_state: LapState, image: &RgbImage) state.frames_without_lap = 0; if state.current_race.is_none() { + state.penalties_this_lap = 0; + let track_hash = get_track_hash(state.config.as_ref(), image); let track_name = state .learned_tracks @@ -89,10 +109,21 @@ fn handle_new_frame(state: &mut AppState, lap_state: LapState, image: &RgbImage) track_hash, track: track_name.unwrap_or_default(), inferred_track, + total_laps: lap_state.total_laps.unwrap_or_default(), + total_positions: lap_state.total_positions.unwrap_or_default(), ..Default::default() }; state.current_race = Some(race); } + + if check_penalty(image, state.config.as_ref()) { + if !state.detecting_penalty { + state.penalties_this_lap += 1; + } + state.detecting_penalty = true; + } else { + state.detecting_penalty = false; + } } else { state.frames_without_lap += 1; } @@ -109,6 +140,8 @@ fn handle_new_frame(state: &mut AppState, lap_state: LapState, image: &RgbImage) merged.lap = Some(lap - 1); } merged.screenshot = Some(to_png_bytes(image)); + merged.penalties = state.penalties_this_lap; + state.penalties_this_lap = 0; if let Some(race) = state.current_race.as_mut() { if let Some(prev_lap) = race.laps.last() { @@ -281,4 +314,21 @@ mod test { && lap_state.lap_time.unwrap() <= Duration::from_secs(50) ); } + + #[test] + fn test_penalty() { + let state = make_test_state(); + let image = image::load_from_memory(include_bytes!("test_data/test-full-1.png")).unwrap(); + analyze_frame(&image.to_rgb8(), &state); + + let penalties = state.lock().unwrap().penalties_this_lap; + assert_eq!(0, penalties); + + let image = + image::load_from_memory(include_bytes!("test_data/test-full-penalty.png")).unwrap(); + analyze_frame(&image.to_rgb8(), &state); + + let penalties = state.lock().unwrap().penalties_this_lap; + assert_eq!(1, penalties); + } } diff --git a/src/config.rs b/src/config.rs index 3140029..cad16ae 100644 --- a/src/config.rs +++ b/src/config.rs @@ -19,6 +19,8 @@ fn default_font_scale() -> f32 { pub struct Config { pub ocr_regions: Vec, pub track_region: Option, + pub penalty_orange_region: Option, + pub penalty_orange_color: Option<(u8, u8, u8)>, #[serde(default = "default_ocr_interval_ms")] pub ocr_interval_ms: u64, @@ -40,10 +42,7 @@ impl Config { load_config_or_make_default("config.json", include_str!("configs/config.default.json")) } - pub fn update_and_save( - self: &mut Arc, - update_fn: F, - ) -> Result<()> { + pub fn update_and_save(self: &mut Arc, update_fn: F) -> Result<()> { let mut config = (**self).clone(); update_fn(&mut config); save_json_config("config.json", &config)?; diff --git a/src/configs/config.default.json b/src/configs/config.default.json index 3b82065..47b2f6d 100644 --- a/src/configs/config.default.json +++ b/src/configs/config.default.json @@ -69,6 +69,16 @@ "threshold": 0.85, "use_ocr_cache": null }, + "penalty_orange_region": { + "name": "penalty", + "x": 993, + "y": 120, + "width": 50, + "height": 30, + "threshold": 0.90, + "use_ocr_cache": null + }, + "penalty_orange_color": [255, 177, 0], "ocr_interval_ms": 500, "track_recognition_threshold": 10, "dump_frame_fraction": null, diff --git a/src/image_processing.rs b/src/image_processing.rs index c465577..db8a1e8 100644 --- a/src/image_processing.rs +++ b/src/image_processing.rs @@ -77,7 +77,7 @@ pub fn check_target_color_fraction( } } } - ((region.height() * region.width()) as f64) / (color_area as f64) + (color_area as f64) / ((region.height() * region.width()) as f64) } pub fn to_png_bytes(image: &RgbImage) -> Vec { diff --git a/src/main.rs b/src/main.rs index 14404a9..18cffe8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -178,6 +178,8 @@ fn show_race_state( config: &Config, ocr_db: &OcrDatabase, ) { + ui.label(format!("Total laps: {}", race.total_laps)); + ui.label(format!("Total competitors: {}", race.total_positions)); egui::Grid::new(format!("race:{}", race_name)).show(ui, |ui| { ui.label("Lap"); ui.label("Position"); @@ -187,6 +189,7 @@ fn show_race_state( ui.label("Health"); ui.label("Gas"); ui.label("Tyres"); + ui.label("Penalties"); ui.end_row(); let mut prev_lap: Option<&LapState> = None; let fastest_lap = race.fastest_lap(); @@ -211,7 +214,7 @@ fn show_race_state( ); show_optional_usize_with_diff( ui, - Color32::GRAY, + Color32::LIGHT_BLUE, &prev_lap.and_then(|p| p.gas), &lap.gas, ); @@ -221,6 +224,7 @@ fn show_race_state( &prev_lap.and_then(|p| p.tyres), &lap.tyres, ); + ui.colored_label(Color32::from_rgb(255, 177, 0), lap.penalties.to_string()); if lap.striked { ui.colored_label(Color32::RED, "Striked from the record"); @@ -336,55 +340,78 @@ impl eframe::App for AppUi { self.ui_state.debug_lap = None; } - egui::SidePanel::left("frame").min_width(180.).show(ctx, |ui| { - if let Some(frame) = &state.last_frame { - ui.heading("Race data"); - ui.label(format!("Lap: {}/{}", frame.lap.unwrap_or(0), frame.total_laps.unwrap_or(0))); - ui.label(format!("Position: {}/{}", frame.position.unwrap_or(0), frame.total_positions.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() - )); - - if let Some(race) = &state.current_race { + egui::SidePanel::left("frame") + .min_width(180.) + .show(ctx, |ui| { + if let Some(frame) = &state.last_frame { + ui.heading("Race data"); + ui.label(format!( + "Lap: {}/{}", + frame.lap.unwrap_or(0), + frame.total_laps.unwrap_or(0) + )); + ui.label(format!( + "Position: {}/{}", + frame.position.unwrap_or(0), + frame.total_positions.unwrap_or(0) + )); ui.separator(); - ui.heading("Strategy"); - if let Some(tyre_wear) = race.tyre_wear() { - ui.label(&format!("Median Tyre Wear: {}", tyre_wear)); - if let Some(tyres) = frame.tyres { - ui.label(&format!( - "Out of tires in {:.1} lap(s)", - (tyres as f64) / (tyre_wear as f64) - )); + ui.colored_label( + Color32::RED, + format!("Health: {}/100", frame.health.unwrap_or(0)), + ); + ui.colored_label( + Color32::LIGHT_BLUE, + format!("Gas: {}/100", frame.gas.unwrap_or(0)), + ); + ui.colored_label( + Color32::GREEN, + format!("Tyres: {}/100", frame.tyres.unwrap_or(0)), + ); + ui.colored_label( + Color32::from_rgb(255, 177, 0), + format!("Penalties: {}", state.penalties_this_lap), + ); + 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() + )); + + if let Some(race) = &state.current_race { + ui.separator(); + ui.heading("Strategy"); + if let Some(tyre_wear) = race.tyre_wear() { + ui.label(&format!("Median Tyre Wear: {}", tyre_wear)); + if let Some(tyres) = frame.tyres { + ui.label(&format!( + "Out of tires in {:.1} lap(s)", + (tyres as f64) / (tyre_wear as f64) + )); + } } - } - if let Some(gas_wear) = race.gas_per_lap() { - ui.label(&format!("Median Gas Wear: {}", gas_wear)); - if let Some(gas) = frame.gas { - ui.label(&format!( - "Out of gas in {:.1} lap(s)", - (gas as f64) / (gas_wear as f64) - )); + if let Some(gas_wear) = race.gas_per_lap() { + ui.label(&format!("Median Gas Wear: {}", gas_wear)); + if let Some(gas) = frame.gas { + ui.label(&format!( + "Out of gas in {:.1} lap(s)", + (gas as f64) / (gas_wear as f64) + )); + } } } } - } - }); + }); egui::CentralPanel::default().show(ctx, |ui| { egui::menu::bar(ui, |ui| { egui::menu::menu_button(ui, "Control", |ui| { diff --git a/src/ocr.rs b/src/ocr.rs index 1f92981..ade8037 100644 --- a/src/ocr.rs +++ b/src/ocr.rs @@ -12,14 +12,14 @@ struct BoundingBox { } fn is_dark_pixel(image: &RgbImage, x: u32, y: u32) -> bool { - let [r, g, b] = image.get_pixel(x, y).0; - r < 100 && g < 100 && b < 100 + let [r, g, b] = image.get_pixel(x, y).0; + r < 100 && g < 100 && b < 100 } fn column_has_any_dark(image: &RgbImage, x: u32) -> bool { for y in 0..image.height() { if is_dark_pixel(image, x, y) { - return true + return true; } } false @@ -51,8 +51,8 @@ fn find_vertical_cutoff_from_right(image: &RgbImage) -> Option<(u32, u32)> { } } } - cutoffs.sort(); - cutoffs.into_iter().skip(5).nth(0) + cutoffs.sort_unstable(); + cutoffs.into_iter().nth(5) } fn get_character_bounding_boxes(image: &RgbImage) -> Vec { diff --git a/src/state.rs b/src/state.rs index 8a030c4..5a89d26 100644 --- a/src/state.rs +++ b/src/state.rs @@ -28,6 +28,8 @@ pub struct LapState { pub striked: bool, pub screenshot: Option>, pub debug: bool, + + pub penalties: usize, } fn parse_duration(time: &str) -> Option { @@ -61,7 +63,7 @@ fn parse_with_slash(v: Option<&String>) -> (Option, Option) { Some(v) => v, None => return (None, None), }; - match v.split_once("/") { + match v.split_once('/') { Some((a, b)) => (a.parse().ok(), b.parse().ok()), None => (v.parse().ok(), None), } @@ -105,6 +107,9 @@ pub struct RaceState { pub laps: Vec, pub last_lap_record_time: Option, + pub total_laps: usize, + pub total_positions: usize, + pub screencap: Option, pub track_hash: Option, @@ -157,6 +162,9 @@ pub struct AppState { pub buffered_frames: VecDeque, + pub penalties_this_lap: usize, + pub detecting_penalty: bool, + pub frames_without_lap: usize, pub current_race: Option, pub past_races: VecDeque, diff --git a/src/test_data/test-full-penalty.png b/src/test_data/test-full-penalty.png new file mode 100644 index 0000000..2073046 Binary files /dev/null and b/src/test_data/test-full-penalty.png differ diff --git a/src/training_ui.rs b/src/training_ui.rs index 348d2eb..89575e0 100644 --- a/src/training_ui.rs +++ b/src/training_ui.rs @@ -98,8 +98,12 @@ impl eframe::App for TrainingUi { egui::CentralPanel::default().show(ctx, |ui| { egui::ScrollArea::vertical().show(ui, |ui| { let training_images_len = self.training_images.len(); - ui.heading(format!("{}/{}", self.current_image_index, training_images_len)); - let current_image = &mut self.training_images[self.current_image_index % training_images_len]; + ui.heading(format!( + "{}/{}", + self.current_image_index, training_images_len + )); + let current_image = + &mut self.training_images[self.current_image_index % training_images_len]; if ui.button("Skip").clicked() { self.current_image_index += 1; } diff --git a/suggestions.md b/suggestions.md index a98eef9..ccbe0be 100644 --- a/suggestions.md +++ b/suggestions.md @@ -1,13 +1,14 @@ - Recognize - - Penalties - Pit stops - Best time from other racers - Editable lap stats - Autocomplete for car/track - Detect car name from load screen - TTS for pit strategies - - Detect position, max laps - - Re-do debug/learn functionality + - Show delta positions + - [DONE] Track penalties + - [DONE] Re-do debug/learn functionality + - [DONE] Detect position, max laps - [DONE] ComboBox for car/track - [DONE] Global best time not current best - [DONE] Don't store uncompressed data for lap debugging \ No newline at end of file