// {{{ Imports use std::fmt::Display; use anyhow::{anyhow, bail}; use image::imageops::FilterType; use image::{DynamicImage, GenericImageView}; use num::integer::Roots; use poise::serenity_prelude::{CreateAttachment, CreateEmbed}; use poise::CreateReply; use crate::arcaea::chart::{Chart, Difficulty, Song, DIFFICULTY_MENU_PIXEL_COLORS}; use crate::arcaea::jacket::IMAGE_VEC_DIM; use crate::arcaea::score::Score; use crate::bitmap::{Color, Rect}; use crate::commands::discord::MessageContext; use crate::context::{Error, UserContext}; use crate::levenshtein::edit_distance; use crate::logs::debug_image_log; use crate::recognition::ui::{ ScoreScreenRect, SongSelectRect, UIMeasurementRect, UIMeasurementRect::*, }; use crate::transform::rotate; // }}} #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ScoreKind { SongSelect, ScoreScreen, } /// Caches a byte vector in order to prevent reallocation #[derive(Debug, Clone, Default)] pub struct ImageAnalyzer { /// cached byte array pub bytes: Vec<u8>, /// Last rect used to crop something last_rect: Option<(UIMeasurementRect, Rect)>, } impl ImageAnalyzer { /// Similar to reinitializing this, but without deallocating memory #[inline] pub fn clear(&mut self) { self.bytes.clear(); self.last_rect = None; } // {{{ Crop #[inline] pub fn crop(&mut self, image: &DynamicImage, rect: Rect) -> DynamicImage { image.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height) } #[inline] pub fn interp_crop( &mut self, ctx: &UserContext, image: &DynamicImage, ui_rect: UIMeasurementRect, ) -> Result<DynamicImage, Error> { let rect = ctx.ui_measurements.interpolate(ui_rect, image)?; self.last_rect = Some((ui_rect, rect)); let result = self.crop(image, rect); debug_image_log(&result); Ok(result) } #[inline] pub fn interp_crop_resize( &mut self, ctx: &UserContext, image: &DynamicImage, ui_rect: UIMeasurementRect, size: (u32, u32), ) -> Result<DynamicImage, Error> { let rect = ctx.ui_measurements.interpolate(ui_rect, image)?; self.last_rect = Some((ui_rect, rect)); let result = self.crop(image, rect); let result = result.resize(size.0, size.1, FilterType::Nearest); debug_image_log(&result); Ok(result) } // }}} // {{{ Error handling pub async fn send_discord_error( &mut self, ctx: &mut impl MessageContext, image: &DynamicImage, filename: &str, err: impl Display, ) -> Result<(), Error> { let mut embed = CreateEmbed::default().description(format!( "Nerdy info ``` {} ```", err )); if let Some((ui_rect, rect)) = self.last_rect { self.crop(image, rect); let bytes = std::mem::take(&mut self.bytes); let error_attachement = CreateAttachment::bytes(bytes, filename); embed = embed.attachment(filename).title(format!( "An error occurred, around the time I was extracting data for {ui_rect:?}" )); ctx.send( CreateReply::default() .embed(embed) .attachment(error_attachement), ) .await?; } else { embed = embed.title("An error occurred"); ctx.send(CreateReply::default().embed(embed)).await?; } Ok(()) } // }}} // {{{ Read score pub fn read_score( &mut self, ctx: &UserContext, note_count: Option<u32>, image: &DynamicImage, kind: ScoreKind, ) -> Result<Score, Error> { let image = self.interp_crop( ctx, image, match kind { ScoreKind::SongSelect => SongSelect(SongSelectRect::Score), ScoreKind::ScoreScreen => ScoreScreen(ScoreScreenRect::Score), }, )?; let measurements = match kind { ScoreKind::SongSelect => &ctx.exo_measurements, ScoreKind::ScoreScreen => &ctx.geosans_measurements, }; let result = Score( measurements .recognise(&image, "0123456789'", None, None)? .chars() .filter(|c| *c != '\'') .collect::<String>() .parse()?, ); // Discard scores if it's impossible let valid_analysis = note_count.is_none_or(|note_count| { let (zeta, shinies, score_units) = result.analyse(note_count); 8_000_000 <= zeta.0 && zeta.0 <= 10_000_000 && shinies <= note_count && score_units <= 2 * note_count }); if result.0 <= 10_010_000 && valid_analysis { Ok(result) } else { Err(anyhow!("Score {result} is not vaild")) } } // }}} // {{{ Read difficulty pub fn read_difficulty( &mut self, ctx: &UserContext, image: &DynamicImage, grayscale_image: &DynamicImage, kind: ScoreKind, ) -> Result<Difficulty, Error> { if kind == ScoreKind::SongSelect { let min = DIFFICULTY_MENU_PIXEL_COLORS .iter() .zip(Difficulty::DIFFICULTIES) .min_by_key(|(c, d)| { let rect = ctx .ui_measurements .interpolate( SongSelect(match d { Difficulty::PST => SongSelectRect::Past, Difficulty::PRS => SongSelectRect::Present, Difficulty::FTR => SongSelectRect::Future, _ => SongSelectRect::Beyond, }), image, ) .unwrap(); let image_color = image.get_pixel(rect.x as u32, rect.y as u32); let image_color = Color::from_bytes(image_color.0); let distance = c.distance(image_color); (distance * 10000.0) as u32 }) .unwrap(); return Ok(min.1); } let image = self.interp_crop( ctx, grayscale_image, ScoreScreen(ScoreScreenRect::Difficulty), )?; let text = ctx.kazesawa_bold_measurements.recognise( &image, "PASTPRESENTFUTUREETERNALBEYOND", Some(200), // We can afford to be generous with binarization here None, )?; let difficulty = Difficulty::DIFFICULTIES .iter() .zip(Difficulty::DIFFICULTY_STRINGS) .min_by_key(|(_, difficulty_string)| edit_distance(difficulty_string, &text)) .map(|(difficulty, _)| *difficulty) .ok_or_else(|| anyhow!("Unrecognised difficulty '{}'", text))?; Ok(difficulty) } // }}} // {{{ Read score kind pub fn read_score_kind( &mut self, ctx: &UserContext, image: &DynamicImage, ) -> Result<ScoreKind, Error> { let image = self.interp_crop(ctx, image, PlayKind)?; let text = ctx .kazesawa_measurements .recognise(&image, "ResultSelectaSong ", None, None)?; let result = if edit_distance(&text, "Result") < edit_distance(&text, "SelectaSong") { ScoreKind::ScoreScreen } else { ScoreKind::SongSelect }; Ok(result) } // }}} // {{{ Read jacket pub fn read_jacket<'a>( &mut self, ctx: &'a UserContext, image: &mut DynamicImage, kind: ScoreKind, difficulty: Difficulty, ) -> Result<(&'a Song, &'a Chart), Error> { let rect = ctx.ui_measurements.interpolate( if kind == ScoreKind::ScoreScreen { ScoreScreen(ScoreScreenRect::Jacket) } else { SongSelect(SongSelectRect::Jacket) }, image, )?; let cropped = if kind == ScoreKind::ScoreScreen { image.view(rect.x as u32, rect.y as u32, rect.width, rect.height) } else { let angle = f32::atan2(rect.height as f32, rect.width as f32); let side = rect.height + rect.width; rotate( image, Rect::new(rect.x, rect.y, side, side), (rect.x, rect.y + rect.height as i32), angle, ); let len = (rect.width.pow(2) + rect.height.pow(2)).sqrt(); image.view(rect.x as u32, rect.y as u32 + rect.height, len, len) }; let (distance, song_id) = ctx .jacket_cache .recognise(&*cropped) .ok_or_else(|| anyhow!("Could not recognise jacket"))?; if distance > (IMAGE_VEC_DIM * 3) as f32 { bail!("No known jacket looks like this"); } let (song, chart) = ctx.song_cache.lookup_by_difficulty(song_id, difficulty)?; Ok((song, chart)) } // }}} // {{{ Read distribution pub fn read_distribution( &mut self, ctx: &UserContext, image: &DynamicImage, ) -> Result<(u32, u32, u32), Error> { let mut out = [0; 3]; use ScoreScreenRect::*; static KINDS: [ScoreScreenRect; 3] = [Pure, Far, Lost]; for i in 0..3 { let image = self.interp_crop(ctx, image, ScoreScreen(KINDS[i]))?; out[i] = ctx .kazesawa_bold_measurements // We need to be very strict with binarization here .recognise(&image, "0123456789", Some(30), Some((0.33, 0.85)))? .parse() .unwrap_or(100000); // This will get discarded as making no sense } println!("Ditribution {out:?}"); Ok((out[0], out[1], out[2])) } // }}} // {{{ Read max recall pub fn read_max_recall( &mut self, ctx: &UserContext, image: &DynamicImage, ) -> Result<u32, Error> { let image = self.interp_crop(ctx, image, ScoreScreen(ScoreScreenRect::MaxRecall))?; let max_recall = ctx .exo_measurements // We can afford to be generous with binarization here .recognise(&image, "0123456789", Some(200), None)? .parse()?; Ok(max_recall) } // }}} }