prescientmoon/shimmeringmoon
prescientmoon/shimmeringmoon
1
Fork 0
Code Activity

so many hacks lol

Browse source 
Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
prescientmoon 2024-06-25 20:44:23 +02:00
parent ab12acd916
commit c77f4fdc1d
Signed by: prescientmoon
SSH key fingerprint: SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
3 changed files with 294 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

110
src/commands.rs
View file

@ -7,6 +7,7 @@ use image::imageops::FilterType;
use image::ImageFormat;
use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage};
use poise::{serenity_prelude as serenity, CreateReply};
use sqlx::query;
use tokio::fs::create_dir_all;
// {{{ Help
@ -35,7 +36,7 @@ pub async fn help(
#[poise::command(
prefix_command,
slash_command,
subcommands("magic"),
subcommands("magic", "delete"),
subcommand_required
)]
pub async fn score(_ctx: Context<'_>) -> Result<(), Error> {
@ -110,11 +111,21 @@ pub async fn magic(
// Create cropper and run OCR
let mut cropper = ImageCropper::default();
let edited = CreateReply::default()
.reply(true)
.content(format!("Image {}: reading jacket", i + 1));
handle.edit(ctx, edited).await?;
let song_by_jacket = cropper.read_jacket(ctx.data(), &image);
// This makes OCR more likely to work
let mut ocr_image = image.grayscale().blur(1.);
let edited = CreateReply::default()
.reply(true)
.content(format!("Image {}: reading difficulty", i + 1));
handle.edit(ctx, edited).await?;
let difficulty = match cropper.read_difficulty(&ocr_image) {
// {{{ OCR error handling
Err(err) => {
@ -135,25 +146,10 @@ pub async fn magic(
ocr_image.invert();
let score = match cropper.read_score(&ocr_image) {
// {{{ OCR error handling
Err(err) => {
error_with_image(
ctx,
&cropper.bytes,
&file.filename,
"Could not read score from picture",
&err,
)
.await?;
continue;
}
// }}}
Ok(score) => score,
};
println!("Score: {}", score.0);
let edited = CreateReply::default()
.reply(true)
.content(format!("Image {}: reading title", i + 1));
handle.edit(ctx, edited).await?;
let song_by_name = cropper.read_song(&ocr_image, &ctx.data().song_cache);
let cached_song = match (song_by_jacket, song_by_name) {
@ -166,7 +162,7 @@ pub async fn magic(
"Hey! I could not read the score in the provided picture.",
&format!(
"This can mean one of three things:
1. The image you provided is not that of an Arcaea score
1. The image you provided is *not that of an Arcaea score
2. The image you provided contains a newly added chart that is not in my database yet
3. The image you provided contains character art that covers the chart name. When this happens, I try to make use of the jacket art in order to determine the chart. It is possible that I've never seen the jacket art for this particular song on this particular difficulty. Contact `@prescientmoon` on discord in order to resolve the issue for you & future users playing this chart!
@ -282,7 +278,7 @@ Title error: {}
} // }}}
};
// {{{ Build play
// {{{ Build chart
let song = &cached_song.song;
let chart = cached_song.lookup(difficulty).ok_or_else(|| {
format!(
@ -290,7 +286,32 @@ Title error: {}
difficulty, song.title
)
})?;
// }}}
let edited = CreateReply::default()
.reply(true)
.content(format!("Image {}: reading score", i + 1));
handle.edit(ctx, edited).await?;
let score = match cropper.read_score(Some(chart.note_count), &ocr_image) {
// {{{ OCR error handling
Err(err) => {
error_with_image(
ctx,
&cropper.bytes,
&file.filename,
"Could not read score from picture",
&err,
)
.await?;
continue;
}
// }}}
Ok(score) => score,
};
// {{{ Build play
let play = CreatePlay::new(score, chart, &user)
.with_attachment(file)
.save(&ctx.data())
@ -298,7 +319,7 @@ Title error: {}
// }}}
// }}}
// {{{ Deliver embed
let (embed, attachment) = play.to_embed(&song, &chart).await?;
let (embed, attachment) = play.to_embed(&song, &chart, i).await?;
embeds.push(embed);
if let Some(attachment) = attachment {
attachments.push(attachment);
@ -331,3 +352,46 @@ Title error: {}
Ok(())
}
// }}}
// {{{ Score delete
/// Delete scores, given their IDs.
#[poise::command(prefix_command, slash_command)]
pub async fn delete(
ctx: Context<'_>,
#[description = "Id of score to delete"] ids: Vec<u32>,
) -> Result<(), Error> {
let user = match User::from_context(&ctx).await {
Ok(user) => user,
Err(_) => {
ctx.say("You are not an user in my database, sorry!")
.await?;
return Ok(());
}
};
if ids.len() == 0 {
ctx.reply("Empty ID list provided").await?;
return Ok(());
}
let mut count = 0;
for id in ids {
let res = query!("DELETE FROM plays WHERE id=? AND user_id=?", id, user.id)
.execute(&ctx.data().db)
.await?;
if res.rows_affected() == 0 {
ctx.reply(format!("No play with id {} found", id)).await?;
} else {
count += 1;
}
}
if count > 0 {
ctx.reply(format!("Deleted {} play(s) successfully!", count))
.await?;
}
Ok(())
}
// }}}

28
src/main.rs
View file

@ -10,7 +10,7 @@ mod score;
mod user;
use context::{Error, UserContext};
use poise::serenity_prelude as serenity;
use poise::serenity_prelude::{self as serenity, UserId};
use sqlx::sqlite::SqlitePoolOptions;
use std::{env::var, path::PathBuf, str::FromStr, sync::Arc, time::Duration};
@ -32,8 +32,7 @@ async fn on_error(error: poise::FrameworkError<'_, UserContext, Error>) {
#[tokio::main]
async fn main() {
let data_dir = var("SHIMMERING_DATA_DIR")
.expect("Missing `SHIMMERING_DATA_DIR` env var, see README for more information.");
let data_dir = var("SHIMMERING_DATA_DIR").expect("Missing `SHIMMERING_DATA_DIR` env var");
let pool = SqlitePoolOptions::new()
.connect(&format!("sqlite://{}/db.sqlite", data_dir))
@ -44,13 +43,28 @@ async fn main() {
let options = poise::FrameworkOptions {
commands: vec![commands::help(), commands::score()],
prefix_options: poise::PrefixFrameworkOptions {
prefix: Some("!".into()),
stripped_dynamic_prefix: Some(|_ctx, message, _user_ctx| {
Box::pin(async {
if message.author.bot || Into::<u64>::into(message.author.id) == 1 {
Ok(None)
} else if message.content.starts_with("!") {
Ok(Some(message.content.split_at(1)))
} else if message.guild_id.is_none() {
if message.content.trim().len() == 0 {
Ok(Some(("", "score magic")))
} else {
Ok(Some(("", &message.content[..])))
}
} else {
Ok(None)
}
})
}),
edit_tracker: Some(Arc::new(poise::EditTracker::for_timespan(
Duration::from_secs(3600),
))),
..Default::default()
},
// The global error handler for all error cases that may occur
on_error: |error| Box::pin(on_error(error)),
..Default::default()
};
@ -68,8 +82,8 @@ async fn main() {
.options(options)
.build();
let token = var("SHIMMERING_DISCORD_TOKEN")
.expect("Missing `SHIMMERING_DISCORD_TOKEN` env var, see README for more information.");
let token =
var("SHIMMERING_DISCORD_TOKEN").expect("Missing `SHIMMERING_DISCORD_TOKEN` env var");
let intents =
serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::MESSAGE_CONTENT;

222
src/score.rs
View file

@ -6,8 +6,8 @@ use std::{
};
use edit_distance::edit_distance;
use image::{DynamicImage, GenericImageView};
use num::Rational64;
use image::{imageops::FilterType, DynamicImage, GenericImageView};
use num::{traits::Euclid, Rational64};
use poise::serenity_prelude::{Attachment, AttachmentId, CreateAttachment, CreateEmbed};
use tesseract::{PageSegMode, Tesseract};
@ -198,7 +198,7 @@ fn score_rects() -> &'static [RelativeRect] {
AbsoluteRect::new(1125, 510, 534, 93, ImageDimensions::new(2778, 1284)).to_relative(),
];
process_datapoints(&mut rects);
widen_by(&mut rects, 0.0, 0.01);
widen_by(&mut rects, 0.0, 0.0075);
rects
})
}
@ -270,9 +270,29 @@ pub fn jacket_rects() -> &'static [RelativeRect] {
pub struct Score(pub u32);
impl Score {
// {{{ Score => ζ-Score
/// Returns the zeta score and the number of shinies
pub fn to_zeta(self, note_count: u32) -> (Score, u32) {
// {{{ Score analysis
// {{{ Mini getters
#[inline]
pub fn to_zeta(self, note_count: u32) -> Score {
self.analyse(note_count).0
}
#[inline]
pub fn shinies(self, note_count: u32) -> u32 {
self.analyse(note_count).1
}
#[inline]
pub fn units(self, note_count: u32) -> u32 {
self.analyse(note_count).2
}
// }}}
/// Returns the zeta score, the number of shinies, and the number of score units.
///
/// Pure (and higher) notes reward two score units, far notes reward one, and lost notes reward
/// none.
pub fn analyse(self, note_count: u32) -> (Score, u32, u32) {
// Smallest possible difference between (zeta-)scores
let increment = Rational64::new_raw(5_000_000, note_count as i64).reduced();
let zeta_increment = Rational64::new_raw(2_000_000, note_count as i64).reduced();
@ -286,7 +306,11 @@ impl Score {
let zeta_score_units = Rational64::from_integer(2) * score_units + shinies;
let zeta_score = Score((zeta_increment * zeta_score_units).floor().to_integer() as u32);
(zeta_score, shinies.to_integer() as u32)
(
zeta_score,
shinies.to_integer() as u32,
score_units.to_integer() as u32,
)
}
// }}}
// {{{ Score => Play rating
@ -368,7 +392,7 @@ impl CreatePlay {
user_id: user.id,
discord_attachment_id: None,
score,
zeta_score: score.to_zeta(chart.note_count as u32).0,
zeta_score: score.to_zeta(chart.note_count as u32),
max_recall: None,
far_notes: None,
// TODO: populate these
@ -446,15 +470,61 @@ pub struct Play {
}
impl Play {
// {{{ Play => distribution
pub fn distribution(&self, note_count: u32) -> Option<(u32, u32, u32, u32)> {
if let Some(fars) = self.far_notes {
let (_, shinies, units) = self.score.analyse(note_count);
let (pures, rem) = (units - fars).div_rem_euclid(&2);
if rem == 1 {
println!("The impossible happened: got an invalid amount of far notes!");
return None;
}
let lost = note_count - fars - pures;
let non_max_pures = pures - shinies;
Some((shinies, non_max_pures, fars, lost))
} else {
None
}
}
// }}}
// {{{ Play => status
#[inline]
pub fn status(&self, chart: &Chart) -> Option<String> {
let score = self.score.0;
if score >= 10_000_000 {
let non_max_pures = chart.note_count + 10_000_000 - score;
if non_max_pures == 0 {
Some("MPM".to_string())
} else {
Some(format!("PM (-{})", non_max_pures))
}
} else if let Some(distribution) = self.distribution(chart.note_count) {
// if no lost notes...
if distribution.3 == 0 {
Some(format!("FR (-{}/-{})", distribution.1, distribution.2))
} else {
Some(format!(
"C (-{}/-{}/-{})",
distribution.1, distribution.2, distribution.3
))
}
} else {
None
}
}
// }}}
// {{{ Play to embed
/// Creates a discord embed for this play.
///
/// The `index` variable is only used to create distinct filenames.
pub async fn to_embed(
&self,
song: &Song,
chart: &Chart,
index: usize,
) -> Result<(CreateEmbed, Option<CreateAttachment>), Error> {
let (_, shiny_count) = self.score.to_zeta(chart.note_count);
let attachement_name = format!("{:?}-{:?}.png", song.id, self.score.0);
let attachement_name = format!("{:?}-{:?}-{:?}.png", song.id, self.score.0, index);
let icon_attachement = match &chart.jacket {
Some(path) => Some(
CreateAttachment::file(&tokio::fs::File::open(path).await?, &attachement_name)
@ -463,6 +533,8 @@ impl Play {
None => None,
};
println!("{:?}", self.score.shinies(chart.note_count));
let mut embed = CreateEmbed::default()
.title(format!(
"{} [{:?} {}]",
@ -488,9 +560,13 @@ impl Play {
true,
)
.field("ζ-Grade", self.zeta_score.grade(), true)
.field("Status", "?", true)
.field(
"Status",
self.status(chart).unwrap_or("?".to_string()),
true,
)
.field("Max recall", "?", true)
.field("Breakdown", format!("{}/?/?/?", shiny_count), true);
.field("Id", format!("{}", self.id), true);
if icon_attachement.is_some() {
embed = embed.thumbnail(format!("attachment://{}", &attachement_name));
@ -513,12 +589,13 @@ mod score_tests {
for shiny_count in 0..=note_count {
let score = Score(10000000 + shiny_count);
let zeta_score_units = 4 * (note_count - shiny_count) + 5 * shiny_count;
let (zeta_score, computed_shiny_count) = score.to_zeta(note_count);
let (zeta_score, computed_shiny_count, units) = score.analyse(note_count);
let expected_zeta_score = Rational64::from_integer(zeta_score_units as i64)
* Rational64::new_raw(2000000, note_count as i64).reduced();
assert_eq!(zeta_score, Score(expected_zeta_score.to_integer() as u32));
assert_eq!(computed_shiny_count, shiny_count);
assert_eq!(units, 2 * note_count);
}
}
}
@ -547,9 +624,13 @@ impl ImageCropper {
}
// {{{ Read score
pub fn read_score(&mut self, image: &DynamicImage) -> Result<Score, Error> {
pub fn read_score(
&mut self,
note_count: Option<u32>,
image: &DynamicImage,
) -> Result<Score, Error> {
self.crop_image_to_bytes(
&image,
&image.resize_exact(image.width(), image.height(), FilterType::Nearest),
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), score_rects())
.ok_or_else(|| "Could not find score area in picture")?
.to_absolute(),
@ -560,28 +641,93 @@ impl ImageCropper {
PageSegMode::PsmSingleWord,
PageSegMode::PsmRawLine,
PageSegMode::PsmSingleLine,
PageSegMode::PsmSparseText,
PageSegMode::PsmSingleBlock,
] {
let result = self.read_score_with_mode(mode)?;
results.push(result.0);
// OCR sometimes loses digits
if result.0 < 1_000_000 {
continue;
} else {
return Ok(result);
let result = self.read_score_with_mode(mode, "0123456789'/");
match result {
Ok(result) => {
results.push(result.0);
}
Err(err) => {
println!("OCR score result error: {}", err);
}
}
}
// {{{ Score correction
// The OCR sometimes fails to read "74" with the arcaea font,
// so we try to detect that and fix it
loop {
println!("Attempts: {:?}.", results);
let old_stack_len = results.len();
results = results
.iter()
.flat_map(|result| {
// If the length is correct, we are good to go!
if *result >= 8_000_000 {
vec![*result]
} else {
let mut results = vec![];
for i in [0, 1, 3, 4] {
let d = 10u32.pow(i);
if (*result / d) % 10 == 4 && (*result / d) % 100 != 74 {
let n = d * 10;
results.push((*result / n) * n * 10 + 7 * n + (*result % n));
}
}
results
}
})
.collect();
if old_stack_len == results.len() {
break;
}
}
// }}}
// {{{ Return score if consensus exists
// 1. Discard scores that are known to be impossible
results = results
.into_iter()
.filter(|result| {
8_000_000 <= *result
&& *result <= 10_010_000
&& note_count
.map(|note_count| {
let (zeta, shinies, score_units) = Score(*result).analyse(note_count);
8_000_000 <= zeta.0
&& zeta.0 <= 10_000_000 && shinies <= note_count
&& score_units <= 2 * note_count
})
.unwrap_or(true)
})
.collect();
// 2. Look for consensus
for result in results.iter() {
if results.iter().filter(|e| **e == *result).count() > results.len() / 2 {
return Ok(Score(*result));
}
}
// }}}
results.sort();
results.dedup();
Err(format!(
"Cannot read score, no matter the mode. Attempts: {:?}",
"Cannot read score. Possible values: {:?}.",
results
))?;
unreachable!()
}
fn read_score_with_mode(&mut self, mode: PageSegMode) -> Result<Score, Error> {
fn read_score_with_mode(&mut self, mode: PageSegMode, whitelist: &str) -> Result<Score, Error> {
let mut t = Tesseract::new(None, Some("eng"))?
// .set_variable("classify_bln_numeric_mode", "1'")?
.set_variable("tessedit_char_whitelist", "0123456789'")?
.set_variable("classify_bln_numeric_mode", "1")?
.set_variable("tessedit_char_whitelist", whitelist)?
.set_image_from_mem(&self.bytes)?;
t.set_page_seg_mode(mode);
t = t.recognize()?;
@ -596,10 +742,12 @@ impl ImageCropper {
// ))?;
// }
let text: String = t
.get_text()?
.trim()
let text: String = t.get_text()?.trim().to_string();
println!("Got {}", text);
let text: String = text
.chars()
.map(|char| if char == '/' { '7' } else { char })
.filter(|char| *char != ' ' && *char != '\'')
.collect();
@ -620,13 +768,17 @@ impl ImageCropper {
t.set_page_seg_mode(PageSegMode::PsmRawLine);
t = t.recognize()?;
if t.mean_text_conf() < 10 {
Err("Difficulty text is not readable.")?;
}
let text: &str = &t.get_text()?;
let text = text.trim();
let conf = t.mean_text_conf();
if conf < 10 && conf != 0 {
Err(format!(
"Difficulty text is not readable (confidence = {}, text = {}).",
conf, text
))?;
}
let difficulty = Difficulty::DIFFICULTIES
.iter()
.zip(Difficulty::DIFFICULTY_STRINGS)
@ -671,8 +823,6 @@ impl ImageCropper {
))?;
}
println!("Raw text: {}, confidence: {}", text, t.mean_text_conf());
let lock = cache.lock().map_err(|_| "Poisoned song cache")?;
let cached_song = loop {
let close_enough: Vec<_> = lock