1
Fork 0

AAAAAAAAAAAA, squish images until OCR actually works :3

Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
prescientmoon 2024-08-10 03:08:38 +02:00
parent 51deb8a68a
commit ef940db80d
Signed by: prescientmoon
SSH key fingerprint: SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
10 changed files with 324 additions and 236 deletions

View file

@ -15,6 +15,23 @@
452 153 0 0 Song select — FTR
638 153 0 0 Song select — ETR/BYD
2340 1080 KauanHenzon
228 10 245 57 Play kind
977 418 403 96 Score screen — score
255 401 564 564 Score screen — jacket
350 305 140 34 Score screen — difficulty
1192 783 78 38 Score screen — pures
1192 838 78 38 Score screen — fars
1192 893 78 38 Score screen — losts
549 344 84 36 Score screen — max recall
528 112 1284 85 Score screen — title
84 235 240 48 Song select — score
432 296 676 37 Song select — jacket
83 141 0 0 Song select — PST
251 141 0 0 Song select — PRS
419 141 0 0 Song select — FTR
587 141 0 0 Song select — ETR/BYD
2160 1620 prescientmoon
19 15 273 60 Play kind
841 682 500 94 Score screen — score
@ -31,3 +48,20 @@
199 159 0 0 Song select — PRS
389 159 0 0 Song select — FTR
581 159 0 0 Song select — ETR/BYD
2220 1080 MathNoob
169 16 250 53 Play kind
900 419 439 95 Score screen — score
193 401 565 565 Score screen — jacket
287 304 138 35 Score screen — difficulty
1128 782 86 39 Score screen — pures
1128 837 86 39 Score screen — fars
1128 892 86 39 Score screen — losts
486 345 85 35 Score screen — max recall
346 112 1467 87 Score screen — title
82 233 257 51 Song select — score
393 296 674 38 Song select — jacket
84 142 0 0 Song select — PST
251 142 0 0 Song select — PRS
419 142 0 0 Song select — FTR
593 142 0 0 Song select — ETR/BYD

View file

@ -1,4 +1,4 @@
use std::path::PathBuf;
use std::{num::NonZeroU16, path::PathBuf};
use image::{ImageBuffer, Rgb};
use sqlx::SqlitePool;
@ -132,41 +132,80 @@ impl Chart {
#[derive(Debug, Clone)]
pub struct CachedSong {
pub song: Song,
charts: [Option<Chart>; 5],
chart_ids: [Option<NonZeroU16>; 5],
}
impl CachedSong {
#[inline]
pub fn new(song: Song, charts: [Option<Chart>; 5]) -> Self {
Self { song, charts }
pub fn new(song: Song) -> Self {
Self {
song,
chart_ids: [None; 5],
}
}
}
// }}}
// {{{ Song cache
#[derive(Debug, Clone, Default)]
pub struct SongCache {
pub songs: Vec<Option<CachedSong>>,
pub charts: Vec<Option<Chart>>,
}
impl SongCache {
#[inline]
pub fn lookup_song(&self, id: u32) -> Result<&CachedSong, Error> {
self.songs
.get(id as usize)
.and_then(|i| i.as_ref())
.ok_or_else(|| format!("Could not find song with id {}", id).into())
}
#[inline]
pub fn lookup(&self, difficulty: Difficulty) -> Result<&Chart, Error> {
self.charts
.get(difficulty.to_index())
.and_then(|c| c.as_ref())
.ok_or_else(|| {
format!(
"Could not find difficulty {:?} for song {}",
difficulty, self.song.title
)
.into()
})
pub fn lookup_chart(&self, chart_id: u32) -> Result<(&Song, &Chart), Error> {
let chart = self
.charts
.get(chart_id as usize)
.and_then(|i| i.as_ref())
.ok_or_else(|| format!("Could not find chart with id {}", chart_id))?;
let song = &self.lookup_song(chart.song_id)?.song;
Ok((song, chart))
}
#[inline]
pub fn lookup_mut(&mut self, difficulty: Difficulty) -> Result<&mut Chart, Error> {
pub fn lookup_song_mut(&mut self, id: u32) -> Result<&mut CachedSong, Error> {
self.songs
.get_mut(id as usize)
.and_then(|i| i.as_mut())
.ok_or_else(|| format!("Could not find song with id {}", id).into())
}
#[inline]
pub fn lookup_chart_mut(&mut self, chart_id: u32) -> Result<&mut Chart, Error> {
self.charts
.get_mut(difficulty.to_index())
.and_then(|c| c.as_mut())
.get_mut(chart_id as usize)
.and_then(|i| i.as_mut())
.ok_or_else(|| format!("Could not find chart with id {}", chart_id).into())
}
#[inline]
pub fn lookup_by_difficulty(
&self,
id: u32,
difficulty: Difficulty,
) -> Result<(&Song, &Chart), Error> {
let cached_song = self.lookup_song(id)?;
let chart_id = cached_song.chart_ids[difficulty.to_index()]
.ok_or_else(|| {
format!(
"Could not find difficulty {:?} for song {}",
difficulty, self.song.title
"Cannot find chart {} [{difficulty:?}]",
cached_song.song.title
)
.into()
})
})?
.get() as u32;
let chart = self.lookup_chart(chart_id)?.1;
Ok((&cached_song.song, chart))
}
#[inline]
@ -178,77 +217,13 @@ impl CachedSong {
pub fn charts_mut(&mut self) -> impl Iterator<Item = &mut Chart> {
self.charts.iter_mut().filter_map(|i| i.as_mut())
}
}
// }}}
// {{{ Song cache
#[derive(Debug, Clone, Default)]
pub struct SongCache {
songs: Vec<Option<CachedSong>>,
}
impl SongCache {
#[inline]
pub fn lookup(&self, id: u32) -> Result<&CachedSong, Error> {
self.songs
.get(id as usize)
.and_then(|i| i.as_ref())
.ok_or_else(|| format!("Could not find song with id {}", id).into())
}
#[inline]
pub fn lookup_chart(&self, chart_id: u32) -> Result<(&Song, &Chart), Error> {
self.songs()
.find_map(|item| {
item.charts().find_map(|chart| {
if chart.id == chart_id {
Some((&item.song, chart))
} else {
None
}
})
})
.ok_or_else(|| format!("Could not find chart with id {}", chart_id).into())
}
#[inline]
pub fn lookup_mut(&mut self, id: u32) -> Result<&mut CachedSong, Error> {
self.songs
.get_mut(id as usize)
.and_then(|i| i.as_mut())
.ok_or_else(|| format!("Could not find song with id {}", id).into())
}
#[inline]
pub fn lookup_chart_mut(&mut self, chart_id: u32) -> Result<&mut Chart, Error> {
self.songs_mut()
.find_map(|item| {
item.charts_mut().find_map(|chart| {
if chart.id == chart_id {
Some(chart)
} else {
None
}
})
})
.ok_or_else(|| format!("Could not find chart with id {}", chart_id).into())
}
#[inline]
pub fn songs(&self) -> impl Iterator<Item = &CachedSong> {
self.songs.iter().filter_map(|i| i.as_ref())
}
#[inline]
pub fn songs_mut(&mut self) -> impl Iterator<Item = &mut CachedSong> {
self.songs.iter_mut().filter_map(|i| i.as_mut())
}
// {{{ Populate cache
pub async fn new(pool: &SqlitePool) -> Result<Self, Error> {
let mut result = Self::default();
// {{{ Songs
let songs = sqlx::query!("SELECT * FROM songs").fetch_all(pool).await?;
for song in songs {
let song = Song {
id: song.id as u32,
@ -265,12 +240,11 @@ impl SongCache {
if song_id >= result.songs.len() {
result.songs.resize(song_id + 1, None);
}
let charts = sqlx::query!("SELECT * FROM charts WHERE song_id=?", song.id)
.fetch_all(pool)
.await?;
let mut chart_cache: [Option<_>; 5] = Default::default();
result.songs[song_id] = Some(CachedSong::new(song));
}
// }}}
// {{{ Charts
let charts = sqlx::query!("SELECT * FROM charts").fetch_all(pool).await?;
for chart in charts {
let chart = Chart {
id: chart.id as u32,
@ -284,12 +258,24 @@ impl SongCache {
note_design: chart.note_design,
};
// {{{ Tie chart to song
{
let index = chart.difficulty.to_index();
chart_cache[index] = Some(chart);
result.lookup_song_mut(chart.song_id)?.chart_ids[index] =
Some(NonZeroU16::new(chart.id as u16).unwrap());
}
result.songs[song_id] = Some(CachedSong::new(song, chart_cache));
// }}}
// {{{ Save chart to cache
{
let index = chart.id as usize;
if index >= result.charts.len() {
result.charts.resize(index + 1, None);
}
result.charts[index] = Some(chart);
}
// }}}
}
// }}}
Ok(result)
}

View file

@ -1,4 +1,4 @@
use std::{fs, path::PathBuf, str::FromStr};
use std::{fs, path::PathBuf};
use image::{imageops::FilterType, GenericImageView, Rgba};
use num::Integer;
@ -99,14 +99,12 @@ impl JacketCache {
.into_rgb8(),
));
for song in song_cache.songs_mut() {
for chart in song.charts_mut() {
for chart in song_cache.charts_mut() {
chart.cached_jacket = Some(Jacket {
raw: contents,
bitmap,
});
}
}
Vec::new()
} else {
@ -126,6 +124,7 @@ impl JacketCache {
if !name.ends_with("_256") {
continue;
}
let name = name.strip_suffix("_256").unwrap();
let difficulty = match name {
@ -140,12 +139,16 @@ impl JacketCache {
_ => Err(format!("Unknown jacket suffix {}", name))?,
};
let (song, chart) = guess_chart_name(dir_name, &song_cache, difficulty, true)?;
let (song_id, chart_id) = {
let (song, chart) =
guess_chart_name(dir_name, &song_cache, difficulty, true)?;
(song.id, chart.id)
};
let contents: &'static _ = fs::read(file.path())?.leak();
let image = image::load_from_memory(contents)?;
jacket_vectors.push((song.id, ImageVec::from_image(&image)));
jacket_vectors.push((song_id, ImageVec::from_image(&image)));
let bitmap: &'static _ = Box::leak(Box::new(
image
@ -154,33 +157,9 @@ impl JacketCache {
));
if name == "base" {
let item = song_cache.lookup_mut(song.id).unwrap();
for chart in item.charts_mut() {
let difficulty_num = match chart.difficulty {
Difficulty::PST => "0",
Difficulty::PRS => "1",
Difficulty::FTR => "2",
Difficulty::BYD => "3",
Difficulty::ETR => "4",
};
// We only want to create this path if there's no overwrite for this
// jacket.
let specialized_path = PathBuf::from_str(
&file
.path()
.to_str()
.unwrap()
.replace("base_night", difficulty_num)
.replace("base", difficulty_num),
)
.unwrap();
let dest = chart.jacket_path(data_dir);
if !specialized_path.exists() && !dest.exists() {
std::os::unix::fs::symlink(file.path(), dest)
.expect("Could not symlink jacket");
// Inefficiently iterates over everything, but it's fine for ~1k entries
for chart in song_cache.charts_mut() {
if chart.song_id == song_id && chart.cached_jacket.is_none() {
chart.cached_jacket = Some(Jacket {
raw: contents,
bitmap,
@ -188,9 +167,7 @@ impl JacketCache {
}
}
} else if difficulty.is_some() {
std::os::unix::fs::symlink(file.path(), chart.jacket_path(data_dir))
.expect("Could not symlink jacket");
let chart = song_cache.lookup_chart_mut(chart.id).unwrap();
let chart = song_cache.lookup_chart_mut(chart_id).unwrap();
chart.cached_jacket = Some(Jacket {
raw: contents,
bitmap,

View file

@ -10,6 +10,7 @@ use crate::arcaea::chart::{Chart, Song};
use crate::context::{Error, UserContext};
use crate::user::User;
use super::chart::SongCache;
use super::score::Score;
// {{{ Create play
@ -263,6 +264,7 @@ impl Play {
AND chart_id=?
AND created_at<?
ORDER BY score DESC
LIMIT 1
",
user.id,
chart.id,
@ -351,3 +353,60 @@ impl Play {
// }}}
}
// }}}
// {{{ General functions
pub type PlayCollection<'a> = Vec<(Play, &'a Song, &'a Chart)>;
pub async fn get_b30_plays<'a>(
db: &SqlitePool,
song_cache: &'a SongCache,
user: &User,
) -> Result<Result<PlayCollection<'a>, &'static str>, Error> {
// {{{ DB data fetching
let plays: Vec<DbPlay> = query_as(
"
SELECT id, chart_id, user_id,
created_at, MAX(score) as score, zeta_score,
creation_ptt, creation_zeta_ptt, far_notes, max_recall, discord_attachment_id
FROM plays p
WHERE user_id = ?
GROUP BY chart_id
ORDER BY score DESC
",
)
.bind(user.id)
.fetch_all(db)
.await?;
// }}}
if plays.len() < 30 {
return Ok(Err("Not enough plays found"));
}
// {{{ B30 computation
// NOTE: we reallocate here, although we do not have much of a choice,
// unless we want to be lazy about things
let mut plays: Vec<(Play, &Song, &Chart)> = plays
.into_iter()
.map(|play| {
let play = play.to_play();
let (song, chart) = song_cache.lookup_chart(play.chart_id)?;
Ok((play, song, chart))
})
.collect::<Result<Vec<_>, Error>>()?;
plays.sort_by_key(|(play, _, chart)| -play.score.play_rating(chart.chart_constant));
plays.truncate(30);
// }}}
Ok(Ok(plays))
}
#[inline]
pub fn compute_b30_ptt(plays: &PlayCollection<'_>) -> i32 {
plays
.iter()
.map(|(play, _, chart)| play.score.play_rating(chart.chart_constant))
.sum::<i32>()
/ 30
}
// }}}

View file

@ -3,7 +3,7 @@ use crate::context::{Context, Error};
pub mod chart;
pub mod score;
pub mod stats;
mod utils;
pub mod utils;
// {{{ Help
/// Show this help menu

View file

@ -17,10 +17,11 @@ use poise::{
use sqlx::query_as;
use crate::{
arcaea::chart::{Chart, Song},
arcaea::jacket::BITMAP_IMAGE_SIZE,
arcaea::play::{DbPlay, Play},
arcaea::score::Score,
arcaea::{
jacket::BITMAP_IMAGE_SIZE,
play::{compute_b30_ptt, get_b30_plays, DbPlay},
score::Score,
},
assets::{
get_b30_background, get_count_background, get_difficulty_background, get_grade_background,
get_name_backgound, get_ptt_emblem, get_score_background, get_status_background,
@ -30,6 +31,7 @@ use crate::{
context::{Context, Error},
get_user,
recognition::fuzzy_song_name::guess_song_and_chart,
reply_errors,
user::discord_it_to_discord_user,
};
@ -121,6 +123,7 @@ pub async fn plot(
let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?;
// SAFETY: we limit the amount of plotted plays to 1000.
let plays = query_as!(
DbPlay,
"
@ -128,6 +131,7 @@ pub async fn plot(
WHERE user_id=?
AND chart_id=?
ORDER BY created_at ASC
LIMIT 1000
",
user.id,
chart.id
@ -230,41 +234,13 @@ pub async fn plot(
#[poise::command(prefix_command, slash_command)]
pub async fn b30(ctx: Context<'_>) -> Result<(), Error> {
let user = get_user!(&ctx);
let user_ctx = ctx.data();
let plays = reply_errors!(
ctx,
get_b30_plays(&user_ctx.db, &user_ctx.song_cache, &user).await?
);
let plays: Vec<DbPlay> = query_as(
"
SELECT id, chart_id, user_id,
created_at, MAX(score) as score, zeta_score,
creation_ptt, creation_zeta_ptt, far_notes, max_recall, discord_attachment_id
FROM plays p
WHERE user_id = ?
GROUP BY chart_id
ORDER BY score DESC
",
)
.bind(user.id)
.fetch_all(&ctx.data().db)
.await?;
if plays.len() < 30 {
ctx.reply("Not enough plays found").await?;
return Ok(());
}
// TODO: consider not reallocating everything here
let mut plays: Vec<(Play, &Song, &Chart)> = plays
.into_iter()
.map(|play| {
let play = play.to_play();
// TODO: change the .lookup to perform binary search or something
let (song, chart) = ctx.data().song_cache.lookup_chart(play.chart_id)?;
Ok((play, song, chart))
})
.collect::<Result<Vec<_>, Error>>()?;
plays.sort_by_key(|(play, _, chart)| -play.score.play_rating(chart.chart_constant));
plays.truncate(30);
// {{{ Layout
let mut layout = LayoutManager::default();
let jacket_area = layout.make_box(BITMAP_IMAGE_SIZE, BITMAP_IMAGE_SIZE);
let jacket_with_border = layout.margin_uniform(jacket_area, 3);
@ -284,14 +260,15 @@ pub async fn b30(ctx: Context<'_>) -> Result<(), Error> {
let item_with_margin = layout.margin_xy(item_area, 22, 17);
let (item_grid, item_origins) = layout.repeated_evenly(item_with_margin, (5, 6));
let root = layout.margin_uniform(item_grid, 30);
// layout.normalize(root);
// }}}
// {{{ Rendering prep
let width = layout.width(root);
let height = layout.height(root);
let canvas = BitmapCanvas::new(width, height);
let mut drawer = LayoutDrawer::new(layout, canvas);
// }}}
// {{{ Render background
let bg = get_b30_background();
drawer.blit_rbg(
@ -304,6 +281,7 @@ pub async fn b30(ctx: Context<'_>) -> Result<(), Error> {
bg.dimensions(),
bg.as_raw(),
);
// }}}
for (i, origin) in item_origins.enumerate() {
drawer
@ -610,7 +588,12 @@ pub async fn b30(ctx: Context<'_>) -> Result<(), Error> {
let mut cursor = Cursor::new(&mut out_buffer);
image.write_to(&mut cursor, image::ImageFormat::Png)?;
let reply = CreateReply::default().attachment(CreateAttachment::bytes(out_buffer, "b30.png"));
let reply = CreateReply::default()
.attachment(CreateAttachment::bytes(out_buffer, "b30.png"))
.content(format!(
"Your ptt is {:.2}",
compute_b30_ptt(&plays) as f32 / 100.0
));
ctx.send(reply).await?;
Ok(())

View file

@ -11,12 +11,18 @@ macro_rules! edit_reply {
#[macro_export]
macro_rules! get_user {
($ctx:expr) => {
match crate::user::User::from_context($ctx).await {
Ok(user) => user,
Err(_) => {
$ctx.say("You are not an user in my database, sorry!")
.await?;
($ctx:expr) => {{
crate::reply_errors!($ctx, crate::user::User::from_context($ctx).await)
}};
}
#[macro_export]
macro_rules! reply_errors {
($ctx:expr, $value:expr) => {
match $value {
Ok(v) => v,
Err(err) => {
$ctx.reply(format!("{err}")).await?;
return Ok(());
}
}

View file

@ -58,15 +58,15 @@ pub fn guess_chart_name<'a>(
let (song, chart) = loop {
let mut close_enough: Vec<_> = cache
.songs()
.filter_map(|item| {
let song = &item.song;
let chart = if let Some(difficulty) = difficulty {
item.lookup(difficulty).ok()?
} else {
item.charts().next()?
};
.charts()
.filter_map(|chart| {
if let Some(difficulty) = difficulty
&& chart.difficulty != difficulty
{
return None;
}
let song = &cache.lookup_song(chart.song_id).ok()?.song;
let song_title = &song.lowercase_title;
distance_vec.clear();

View file

@ -4,7 +4,8 @@ use std::str::FromStr;
use std::{env, fs};
use hypertesseract::{PageSegMode, Tesseract};
use image::{DynamicImage, GenericImageView};
use image::imageops::{resize, FilterType};
use image::{DynamicImage, GenericImageView, RgbaImage};
use image::{ImageBuffer, Rgba};
use num::integer::Roots;
use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage, Timestamp};
@ -46,26 +47,27 @@ impl ImageAnalyzer {
}
// {{{ Crop
pub fn crop_image_to_bytes(&mut self, image: &DynamicImage, rect: Rect) -> Result<(), Error> {
#[inline]
fn should_save_debug_images() -> bool {
env::var("SHIMMERING_DEBUG_IMGS")
.map(|s| s == "1")
.unwrap_or(false)
}
fn save_image(&mut self, image: &RgbaImage) -> Result<(), Error> {
self.clear();
let image = image.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height);
let mut cursor = Cursor::new(&mut self.bytes);
image.write_to(&mut cursor, image::ImageFormat::Png)?;
if Self::should_save_debug_images() {
fs::write(format!("./logs/{}.png", Timestamp::now()), &self.bytes)?;
}
Ok(())
}
#[inline]
pub fn crop(&mut self, image: &DynamicImage, rect: Rect) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
if env::var("SHIMMERING_DEBUG_IMGS")
.map(|s| s == "1")
.unwrap_or(false)
{
self.crop_image_to_bytes(image, rect).unwrap();
}
image
.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height)
.to_rgba8()
@ -80,7 +82,35 @@ impl ImageAnalyzer {
) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>, Error> {
let rect = ctx.ui_measurements.interpolate(ui_rect, image)?;
self.last_rect = Some((ui_rect, rect));
Ok(self.crop(image, rect))
let result = self.crop(image, rect);
if Self::should_save_debug_images() {
self.save_image(&result).unwrap();
}
Ok(result)
}
#[inline]
pub fn interp_crop_resize(
&mut self,
ctx: &UserContext,
image: &DynamicImage,
ui_rect: UIMeasurementRect,
size: impl FnOnce(Rect) -> (u32, u32),
) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>, Error> {
let rect = ctx.ui_measurements.interpolate(ui_rect, image)?;
let size = size(rect);
self.last_rect = Some((ui_rect, rect));
let result = self.crop(image, rect);
let result = resize(&result, size.0, size.1, FilterType::Nearest);
if Self::should_save_debug_images() {
self.save_image(&result).unwrap();
}
Ok(result)
}
// }}}
// {{{ Error handling
@ -100,7 +130,8 @@ impl ImageAnalyzer {
));
if let Some((ui_rect, rect)) = self.last_rect {
self.crop_image_to_bytes(image, rect)?;
let cropped = self.crop(image, rect);
self.save_image(&cropped)?;
let bytes = std::mem::take(&mut self.bytes);
let error_attachement = CreateAttachment::bytes(bytes, filename);
@ -131,13 +162,26 @@ impl ImageAnalyzer {
image: &DynamicImage,
kind: ScoreKind,
) -> Result<Vec<Score>, Error> {
let image = self.interp_crop(
// yes, this was painfully hand-picked
let desired_height = 100;
let x_scaling_factor = match kind {
ScoreKind::SongSelect => 1.0,
ScoreKind::ScoreScreen => 0.66,
};
let image = self.interp_crop_resize(
ctx,
image,
if kind == ScoreKind::ScoreScreen {
ScoreScreen(ScoreScreenRect::Score)
} else {
SongSelect(SongSelectRect::Score)
match kind {
ScoreKind::SongSelect => SongSelect(SongSelectRect::Score),
ScoreKind::ScoreScreen => ScoreScreen(ScoreScreenRect::Score),
},
|rect| {
(
(rect.width as f32 * desired_height as f32 / rect.height as f32
* x_scaling_factor) as u32,
desired_height,
)
},
)?;
@ -423,11 +467,9 @@ impl ImageAnalyzer {
Err("No known jacket looks like this")?;
}
let item = ctx.song_cache.lookup(*song_id)?;
let chart = item.lookup(difficulty)?;
let (song, chart) = ctx.song_cache.lookup_by_difficulty(*song_id, difficulty)?;
// NOTE: this will reallocate a few strings, but it is what it is
Ok((&item.song, chart))
Ok((song, chart))
}
// }}}
// {{{ Read distribution

View file

@ -16,7 +16,8 @@ impl User {
let id = ctx.author().id.get().to_string();
let user = sqlx::query!("SELECT * FROM users WHERE discord_id = ?", id)
.fetch_one(&ctx.data().db)
.await?;
.await
.map_err(|_| "You are not an user in my database, sowwy ^~^")?;
Ok(User {
id: user.id as u32,