prescientmoon/shimmeringmoon
1
Fork 0
Code Activity

So much work, my god

Browse source 
Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
prescientmoon 2024-06-23 02:51:50 +02:00
parent 5cfeff4e14
commit 5c4bfa25c9
Signed by: prescientmoon
SSH key fingerprint: SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
9 changed files with 460 additions and 183 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
data/jackets.csv
View file

@ -1,4 +1,3 @@
filename,song_id
grievous-lady,7
einherjar-joker,14
einherjar-joker-byd,14
Grievous Lady,grievous-lady,
Einherjar Joker,einherjar-joker,
Einherjar Joker,einherjar-joker-byd,BYD

1 filename Grievous Lady song_id grievous-lady
2 grievous-lady Einherjar Joker 7 einherjar-joker
3 einherjar-joker Einherjar Joker 14 einherjar-joker-byd BYD
einherjar-joker-byd 14

1
data/shorthands.csv Normal file
View file

@ -0,0 +1 @@
Misdeed -la bonté de Dieu et l'origine du mal-, Misdeed -la bonte de Dieu et lorigine du mal-
1 Misdeed -la bonté de Dieu et l'origine du mal- Misdeed -la bonte de Dieu et lorigine du mal-

1
schema.sql
View file

@ -19,6 +19,7 @@ CREATE TABLE IF NOT EXISTS songs (
CREATE TABLE IF NOT EXISTS charts (
id INTEGER NOT NULL PRIMARY KEY,
song_id INTEGER NOT NULL,
jacket TEXT,
difficulty TEXT NOT NULL CHECK (difficulty IN ('PST','PRS','FTR','ETR','BYD')),
level TEXT NOT NULL,

97
scripts/main.py
View file

@ -14,54 +14,79 @@ conn = sqlite3.connect(db_path)
# {{{ Import songs
def import_charts_from_csv(input_file):
with open(input_file, mode="r") as file:
chart_count = 0
songs = dict()
def import_charts_from_csv():
chart_count = 0
songs = dict()
with open(data_dir + "/charts.csv", mode="r") as file:
for row in csv.reader(file):
if len(row) > 0:
chart_count += 1
[title, difficulty, level, cc, _, note_count, _, _, _] = row
if songs.get(title) is None:
songs[title] = []
songs[title].append((difficulty, level, cc, note_count))
songs[title] = {"charts": [], "shorthand": None}
songs[title]["charts"].append([difficulty, level, cc, note_count, None])
for title, charts in songs.items():
artist = None
with open(data_dir + "/jackets.csv", mode="r") as file:
for row in csv.reader(file):
if len(row) > 0:
[title, jacket, difficulty] = row
if difficulty.strip() != "":
changed = 0
if title.startswith("Quon"):
artist = title[6:-1]
title = "Quon"
for i in range(len(songs[title]["charts"])):
if songs[title]["charts"][i][0] == difficulty:
songs[title]["charts"][i][4] = jacket
changed += 1
row = conn.execute(
if changed == 0:
raise f"Nothing changed for chart {title} [{difficulty}]"
else:
for i in range(len(songs[title]["charts"])):
songs[title]["charts"][i][4] = jacket
with open(data_dir + "/shorthands.csv", mode="r") as file:
for row in csv.reader(file):
if len(row) > 0:
[title, shorthand] = row
songs[title]["shorthand"] = shorthand
for title, entry in songs.items():
artist = None
if title.startswith("Quon"):
artist = title[6:-1]
title = "Quon"
row = conn.execute(
"""
INSERT INTO songs(title,artist,ocr_alias)
VALUES (?,?,?)
RETURNING id
""",
(title, artist, entry.get("shorthand")),
).fetchone()
song_id = row[0]
for difficulty, level, cc, note_count, jacket in entry["charts"]:
conn.execute(
"""
INSERT INTO songs(title,artist)
VALUES (?,?)
RETURNING id
""",
(title, artist),
).fetchone()
song_id = row[0]
for difficulty, level, cc, note_count in charts:
conn.execute(
"""
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant)
VALUES(?,?,?,?,?)
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant, jacket)
VALUES(?,?,?,?,?, ?)
""",
(
song_id,
difficulty,
level,
int(note_count.replace(",", "").replace(".", "")),
int(float(cc) * 100),
),
)
(
song_id,
difficulty,
level,
int(note_count.replace(",", "").replace(".", "")),
int(float(cc) * 100),
jacket,
),
)
conn.commit()
conn.commit()
print(f"Imported {chart_count} charts and {len(songs)} songs")
print(f"Imported {chart_count} charts and {len(songs)} songs")
# }}}
@ -70,4 +95,4 @@ command = sys.argv[1]
subcommand = sys.argv[2]
if command == "import" and subcommand == "charts":
import_charts_from_csv(sys.argv[3])
import_charts_from_csv()

43
src/chart.rs
View file

@ -1,3 +1,5 @@
use std::path::PathBuf;
use sqlx::{prelude::FromRow, SqlitePool};
use crate::context::Error;
@ -46,6 +48,13 @@ pub struct Song {
pub ocr_alias: Option<String>,
pub artist: Option<String>,
}
impl Song {
#[inline]
pub fn ocr_string(&self) -> &str {
(&self.ocr_alias).as_ref().unwrap_or(&self.title)
}
}
// }}}
// {{{ Chart
#[derive(Debug, Clone, FromRow)]
@ -58,9 +67,11 @@ pub struct Chart {
pub note_count: u32,
pub chart_constant: u32,
pub jacket: Option<PathBuf>,
}
// }}}
// {{{ Cache
// {{{ Cached song
#[derive(Debug, Clone)]
pub struct CachedSong {
pub song: Song,
@ -79,8 +90,21 @@ impl CachedSong {
.get(difficulty.to_index())
.and_then(|c| c.as_ref())
}
}
#[inline]
pub fn lookup_mut(&mut self, difficulty: Difficulty) -> Option<&mut Chart> {
self.charts
.get_mut(difficulty.to_index())
.and_then(|c| c.as_mut())
}
#[inline]
pub fn charts(&self) -> impl Iterator<Item = &Chart> {
self.charts.iter().filter_map(|i| i.as_ref())
}
}
// }}}
// {{{ Song cache
#[derive(Debug, Clone, Default)]
pub struct SongCache {
songs: Vec<Option<CachedSong>>,
@ -92,8 +116,18 @@ impl SongCache {
self.songs.get(id as usize).and_then(|i| i.as_ref())
}
#[inline]
pub fn lookup_mut(&mut self, id: u32) -> Option<&mut CachedSong> {
self.songs.get_mut(id as usize).and_then(|i| i.as_mut())
}
#[inline]
pub fn songs(&self) -> impl Iterator<Item = &CachedSong> {
self.songs.iter().filter_map(|i| i.as_ref())
}
// {{{ Populate cache
pub async fn new(pool: &SqlitePool) -> Result<Self, Error> {
pub async fn new(data_dir: &PathBuf, pool: &SqlitePool) -> Result<Self, Error> {
let mut result = Self::default();
let songs = sqlx::query!("SELECT * FROM songs").fetch_all(pool).await?;
@ -125,6 +159,9 @@ impl SongCache {
level: chart.level,
chart_constant: chart.chart_constant as u32,
note_count: chart.note_count as u32,
jacket: chart
.jacket
.map(|jacket| data_dir.join("jackets").join(format!("{}.png", jacket))),
};
let index = chart.difficulty.to_index();

235
src/commands.rs
View file

@ -1,12 +1,15 @@
use std::fmt::Display;
use crate::context::{Context, Error};
use crate::score::{CreatePlay, ImageCropper};
use crate::score::{jacket_rects, CreatePlay, ImageCropper, ImageDimensions, RelativeRect};
use crate::user::User;
use image::imageops::FilterType;
use image::ImageFormat;
use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage};
use poise::{serenity_prelude as serenity, CreateReply};
use tokio::fs::create_dir_all;
// {{{ Help
/// Show this help menu
#[poise::command(prefix_command, track_edits, slash_command)]
pub async fn help(
@ -26,7 +29,8 @@ pub async fn help(
.await?;
Ok(())
}
// }}}
// {{{ Score
/// Score management
#[poise::command(
prefix_command,
@ -37,7 +41,8 @@ pub async fn help(
pub async fn score(_ctx: Context<'_>) -> Result<(), Error> {
Ok(())
}
// }}}
// {{{ Score magic
// {{{ Send error embed with image
async fn error_with_image(
ctx: Context<'_>,
@ -77,96 +82,40 @@ pub async fn magic(
}
};
println!("Handling command from user {:?}", user.discord_id);
if files.len() == 0 {
ctx.reply("No images found attached to message").await?;
} else {
let mut embeds: Vec<CreateEmbed> = vec![];
let mut attachements: Vec<CreateAttachment> = vec![];
let mut attachments: Vec<CreateAttachment> = vec![];
let handle = ctx
.reply(format!("Processed 0/{} scores", files.len()))
.await?;
for (i, file) in files.iter().enumerate() {
if let Some(_) = file.dimensions() {
// {{{ Image pre-processing
// Download image and guess it's format
let bytes = file.download().await?;
let format = image::guess_format(&bytes)?;
// Image pre-processing
let image = image::load_from_memory_with_format(&bytes, format)?.resize(
1024,
1024,
FilterType::Nearest,
);
// // {{{ Table experiment
// let table_format = FormatBuilder::new()
// .separators(
// &[LinePosition::Title],
// LineSeparator::new('─', '┬', '┌', '┐'),
// )
// .padding(1, 1)
// .build();
// let mut table = Table::new();
// table.set_format(table_format);
// table.set_titles(row!["Chart", "Level", "Score", "Rating"]);
// table.add_row(row!["Quon", "BYD 10", "10000807", "12.3 (-132)"]);
// table.add_row(row!["Monochrome princess", "FTR 9+", " 9380807", "10.2"]);
// table.add_row(row!["Grievous lady", "FTR 11", " 9286787", "11.2"]);
// table.add_row(row!["Fracture ray", "FTR 11", " 8990891", "11.0"]);
// table.add_row(row!["Shades of Light", "FTR 9+", "10000976", " 9.3 (-13)"]);
// ctx.say(format!("```\n{}\n```", table.to_string())).await?;
// // }}}
// // {{{ Embed experiment
// let icon_attachement = CreateAttachment::file(
// &tokio::fs::File::open("./data/jackets/grievous.png").await?,
// "grievous.png",
// )
// .await?;
// let msg = CreateMessage::default().embed(
// CreateEmbed::default()
// .title("Grievous lady [FTR 11]")
// .thumbnail("attachment://grievous.png")
// .field("Score", "998302 (+8973)", true)
// .field("Rating", "12.2 (+.6)", true)
// .field("Grade", "EX+", true)
// .field("ζ-Score", "982108 (+347)", true)
// .field("ζ-Rating", "11.5 (+.45)", true)
// .field("ζ-Grade", "EX", true)
// .field("Status", "FR (-243F)", true)
// .field("Max recall", "308/1073", true)
// .field("Breakdown", "894/342/243/23", true),
// );
//
// ctx.channel_id()
// .send_files(ctx.http(), [icon_attachement], msg)
// .await?;
// // }}}
// }}}
// {{{ Detection
// Create cropper and run OCR
let mut cropper = ImageCropper::default();
let (jacket, cached_song) = match cropper.read_jacket(ctx.data(), &image) {
// {{{ Jacket recognition error handling
Err(err) => {
error_with_image(
ctx,
&cropper.bytes,
&file.filename,
"Error while detecting jacket",
err,
)
.await?;
let song_by_jacket = cropper.read_jacket(ctx.data(), &image);
continue;
}
// }}}
Ok(j) => j,
};
// This makes OCR more likely to work
let mut ocr_image = image.grayscale().blur(1.);
let mut image = image.grayscale().blur(1.);
let difficulty = match cropper.read_difficulty(&image) {
let difficulty = match cropper.read_difficulty(&ocr_image) {
// {{{ OCR error handling
Err(err) => {
error_with_image(
@ -184,9 +133,9 @@ pub async fn magic(
Ok(d) => d,
};
image.invert();
ocr_image.invert();
let score = match cropper.read_score(&image) {
let score = match cropper.read_score(&ocr_image) {
// {{{ OCR error handling
Err(err) => {
error_with_image(
@ -204,6 +153,136 @@ pub async fn magic(
Ok(score) => score,
};
println!("Score: {}", score.0);
let song_by_name = cropper.read_song(&ocr_image, &ctx.data().song_cache);
let cached_song = match (song_by_jacket, song_by_name) {
// {{{ Both errors
(Err(err_jacket), Err(err_name)) => {
error_with_image(
ctx,
&cropper.bytes,
&file.filename,
"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
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!
Nerdy info:
```
Jacket error: {}
Title error: {}
```" ,
err_jacket, err_name
),
)
.await?;
continue;
}
// }}}
// {{{ Only jacket succeeded
(Ok(by_jacket), Err(err_name)) => {
println!("Could not read name with error: {}", err_name);
by_jacket
}
// }}}
// {{{ Only name succeeded
(Err(err_jacket), Ok(mut by_name)) => {
println!("Could not recognise jacket with error: {}", err_jacket);
// {{{ Find image rect
let rect = RelativeRect::from_aspect_ratio(
ImageDimensions::from_image(&image),
jacket_rects(),
)
.ok_or_else(|| "Could not find jacket area in picture")?
.to_absolute();
// }}}
// {{{ Find chart
let chart = by_name.lookup(difficulty).ok_or_else(|| {
format!(
"Cannot find difficulty {:?} for chart {:?}",
difficulty, by_name.song.title
)
})?;
// }}}
// {{{ Build path
let filename = format!("{}-{}", by_name.song.id, chart.id);
let jacket = format!("user/{}", filename);
let jacket_dir = ctx.data().data_dir.join("jackets/user");
create_dir_all(&jacket_dir).await?;
let jacket_path = jacket_dir.join(format!("{}.png", filename));
// }}}
// {{{ Save image to disk
image
.crop_imm(rect.x, rect.y, rect.width, rect.height)
.save_with_format(&jacket_path, ImageFormat::Png)?;
// }}}
// {{{ Update jacket in db
sqlx::query!(
"UPDATE charts SET jacket=? WHERE song_id=? AND difficulty=?",
jacket,
chart.song_id,
chart.difficulty,
)
.execute(&ctx.data().db)
.await?;
// }}}
// {{{ Aquire and use song cache lock
{
let mut song_cache = ctx
.data()
.song_cache
.lock()
.map_err(|_| "Poisoned song cache")?;
let chart = song_cache
.lookup_mut(by_name.song.id)
.ok_or_else(|| {
format!("Could not find song for id {}", by_name.song.id)
})?
.lookup_mut(difficulty)
.ok_or_else(|| {
format!(
"Could not find difficulty {:?} for song {}",
difficulty, by_name.song.title
)
})?;
if chart.jacket.is_none() {
if let Some(chart) = by_name.lookup_mut(difficulty) {
chart.jacket = Some(jacket_path.clone());
};
chart.jacket = Some(jacket_path);
} else {
println!(
"Jacket not detected for chart {} [{:?}]",
by_name.song.id, difficulty
)
};
}
// }}}
by_name
}
// }}}
// {{{ Both succeeded
(Ok(by_jacket), Ok(by_name)) => {
if by_name.song.id != by_jacket.song.id {
println!(
"Got diverging choices between '{:?}' and '{:?}'",
by_jacket.song.id, by_name.song.id
);
};
by_jacket
} // }}}
};
// {{{ Build play
let song = &cached_song.song;
let chart = cached_song.lookup(difficulty).ok_or_else(|| {
format!(
@ -216,10 +295,15 @@ pub async fn magic(
.with_attachment(file)
.save(&ctx.data())
.await?;
let (embed, attachement) = play.to_embed(&song, &chart, &jacket).await?;
// }}}
// }}}
// {{{ Deliver embed
let (embed, attachment) = play.to_embed(&song, &chart).await?;
embeds.push(embed);
attachements.push(attachement);
if let Some(attachment) = attachment {
attachments.push(attachment);
}
// }}}
} else {
ctx.reply("One of the attached files is not an image!")
.await?;
@ -240,9 +324,10 @@ pub async fn magic(
let msg = CreateMessage::new().embeds(embeds);
ctx.channel_id()
.send_files(ctx.http(), attachements, msg)
.send_files(ctx.http(), attachments, msg)
.await?;
}
Ok(())
}
// }}}

13
src/context.rs
View file

@ -1,4 +1,7 @@
use std::path::PathBuf;
use std::{
path::PathBuf,
sync::{Arc, Mutex},
};
use sqlx::SqlitePool;
@ -12,19 +15,19 @@ pub type Context<'a> = poise::Context<'a, UserContext, Error>;
pub struct UserContext {
pub data_dir: PathBuf,
pub db: SqlitePool,
pub song_cache: SongCache,
pub song_cache: Arc<Mutex<SongCache>>,
pub jacket_cache: JacketCache,
}
impl UserContext {
#[inline]
pub async fn new(data_dir: PathBuf, db: SqlitePool) -> Result<Self, Error> {
let song_cache = SongCache::new(&db).await?;
let jacket_cache = JacketCache::new(&data_dir)?;
let song_cache = SongCache::new(&data_dir, &db).await?;
let jacket_cache = JacketCache::new(&song_cache)?;
Ok(Self {
data_dir,
db,
song_cache,
song_cache: Arc::new(Mutex::new(song_cache)),
jacket_cache,
})
}

43
src/jacket.rs
View file

@ -1,10 +1,10 @@
use std::path::PathBuf;
use std::{collections::HashSet, path::PathBuf};
use image::{GenericImageView, Rgba};
use kd_tree::{KdMap, KdPoint};
use num::Integer;
use crate::context::Error;
use crate::{chart::SongCache, context::Error};
/// How many sub-segments to split each side into
const SPLIT_FACTOR: u32 = 5;
@ -15,12 +15,6 @@ pub struct ImageVec {
pub colors: [f32; IMAGE_VEC_DIM],
}
#[derive(Debug, Clone)]
pub struct Jacket {
pub song_id: u32,
pub path: PathBuf,
}
impl ImageVec {
// {{{ (Image => vector) encoding
fn from_image(image: &impl GenericImageView<Pixel = Rgba<u8>>) -> ImageVec {
@ -76,29 +70,28 @@ impl KdPoint for ImageVec {
}
pub struct JacketCache {
tree: KdMap<ImageVec, Jacket>,
// TODO: make this private
pub tree: KdMap<ImageVec, u32>,
}
impl JacketCache {
// {{{ Generate tree
pub fn new(data_dir: &PathBuf) -> Result<Self, Error> {
let jacket_csv_path = data_dir.join("jackets.csv");
let mut reader = csv::Reader::from_path(jacket_csv_path)?;
// This is a bit inefficient (using a hash set), but only runs once
pub fn new(song_cache: &SongCache) -> Result<Self, Error> {
let mut entries = vec![];
let mut jackets: HashSet<(&PathBuf, u32)> = HashSet::new();
for record in reader.records() {
let record = record?;
let filename = &record[0];
let song_id = u32::from_str_radix(&record[1], 10)?;
let image_path = data_dir.join(format!("jackets/{}.png", filename));
let image = image::io::Reader::open(&image_path)?.decode()?;
let jacket = Jacket {
song_id,
path: image_path,
};
for item in song_cache.songs() {
for chart in item.charts() {
if let Some(jacket) = &chart.jacket {
jackets.insert((jacket, item.song.id));
}
}
}
entries.push((ImageVec::from_image(&image), jacket))
for (path, song_id) in jackets {
let image = image::io::Reader::open(path)?.decode()?;
entries.push((ImageVec::from_image(&image), song_id))
}
let result = Self {
@ -113,7 +106,7 @@ impl JacketCache {
pub fn recognise(
&self,
image: &impl GenericImageView<Pixel = Rgba<u8>>,
) -> Option<(f32, &Jacket)> {
) -> Option<(f32, &u32)> {
self.tree
.nearest(&ImageVec::from_image(image))
.map(|p| (p.squared_distance.sqrt(), &p.item.1))

203
src/score.rs
View file

@ -1,5 +1,9 @@
#![allow(dead_code)]
use std::{fmt::Display, io::Cursor, sync::OnceLock};
use std::{
fmt::Display,
io::Cursor,
sync::{Mutex, OnceLock},
};
use image::{DynamicImage, GenericImageView};
use num::Rational64;
@ -7,9 +11,8 @@ use poise::serenity_prelude::{Attachment, AttachmentId, CreateAttachment, Create
use tesseract::{PageSegMode, Tesseract};
use crate::{
chart::{CachedSong, Chart, Difficulty, Song},
chart::{CachedSong, Chart, Difficulty, Song, SongCache},
context::{Error, UserContext},
jacket::Jacket,
user::User,
};
@ -140,6 +143,7 @@ impl RelativeRect {
}
// }}}
// {{{ Data points
// {{{ Processing
fn process_datapoints(rects: &mut Vec<RelativeRect>) {
rects.sort_by_key(|r| (r.dimensions.aspect_ratio() * 1000.0).floor() as u32);
@ -166,7 +170,7 @@ fn widen_by(rects: &mut Vec<RelativeRect>, x: f32, y: f32) {
rect.height += 2. * y;
}
}
// }}}
// {{{ Score
fn score_rects() -> &'static [RelativeRect] {
static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
@ -183,7 +187,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.005);
widen_by(&mut rects, 0.0, 0.01);
rects
})
}
@ -208,8 +212,29 @@ fn difficulty_rects() -> &'static [RelativeRect] {
})
}
// }}}
// {{{ Chart title
fn title_rects() -> &'static [RelativeRect] {
static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
CELL.get_or_init(|| {
let mut rects: Vec<RelativeRect> = vec![
AbsoluteRect::new(227, 74, 900, 61, ImageDimensions::new(1560, 720)).to_relative(),
AbsoluteRect::new(413, 72, 696, 58, ImageDimensions::new(1600, 720)).to_relative(),
AbsoluteRect::new(484, 148, 1046, 96, ImageDimensions::new(2000, 1200)).to_relative(),
AbsoluteRect::new(438, 324, 1244, 104, ImageDimensions::new(2160, 1620)).to_relative(),
AbsoluteRect::new(216, 336, 1366, 96, ImageDimensions::new(2224, 1668)).to_relative(),
AbsoluteRect::new(634, 116, 1252, 102, ImageDimensions::new(2532, 1170)).to_relative(),
AbsoluteRect::new(586, 222, 1320, 118, ImageDimensions::new(2560, 1600)).to_relative(),
AbsoluteRect::new(348, 417, 1716, 120, ImageDimensions::new(2732, 2048)).to_relative(),
AbsoluteRect::new(760, 128, 1270, 118, ImageDimensions::new(2778, 1284)).to_relative(),
];
process_datapoints(&mut rects);
widen_by(&mut rects, 0.1, 0.0);
rects
})
}
// }}}
// {{{ Jacket
fn jacket_rects() -> &'static [RelativeRect] {
pub fn jacket_rects() -> &'static [RelativeRect] {
static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
CELL.get_or_init(|| {
let mut rects: Vec<RelativeRect> = vec![
@ -238,8 +263,8 @@ impl Score {
/// Returns the zeta score and the number of shinies
pub fn to_zeta(self, note_count: u32) -> (Score, u32) {
// Smallest possible difference between (zeta-)scores
let increment = Rational64::new_raw(5000000, note_count as i64).reduced();
let zeta_increment = Rational64::new_raw(2000000, note_count as i64).reduced();
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();
let score = Rational64::from_integer(self.0 as i64);
let score_units = (score / increment).floor();
@ -257,12 +282,12 @@ impl Score {
#[inline]
pub fn play_rating(self, chart_constant: u32) -> i32 {
chart_constant as i32
+ if self.0 >= 10000000 {
+ if self.0 >= 10_000_000 {
200
} else if self.0 >= 9800000 {
100 + (self.0 as i32 - 9_800_000) / 20_000
} else if self.0 >= 9_800_000 {
100 + (self.0 as i32 - 9_800_000) / 2_000
} else {
(self.0 as i32 - 9_500_000) / 10_000
(self.0 as i32 - 9_500_000) / 3_000
}
}
// }}}
@ -295,7 +320,7 @@ impl Display for Score {
let score = self.0;
write!(
f,
"{}'{}'{}",
"{}'{:0>3}'{:0>3}",
score / 1000000,
(score / 1000) % 1000,
score % 1000
@ -415,20 +440,23 @@ impl Play {
&self,
song: &Song,
chart: &Chart,
jacket: &Jacket,
) -> Result<(CreateEmbed, CreateAttachment), Error> {
) -> 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 icon_attachement = CreateAttachment::file(
&tokio::fs::File::open(&jacket.path).await?,
&attachement_name,
)
.await?;
let icon_attachement = match &chart.jacket {
Some(path) => Some(
CreateAttachment::file(&tokio::fs::File::open(path).await?, &attachement_name)
.await?,
),
None => None,
};
let embed = CreateEmbed::default()
.title(&song.title)
.thumbnail(format!("attachment://{}", &attachement_name))
let mut embed = CreateEmbed::default()
.title(format!(
"{} [{:?} {}]",
&song.title, chart.difficulty, chart.level
))
.field("Score", format!("{} (+?)", self.score), true)
.field(
"Rating",
@ -453,6 +481,10 @@ impl Play {
.field("Max recall", "?", true)
.field("Breakdown", format!("{}/?/?/?", shiny_count), true);
if icon_attachement.is_some() {
embed = embed.thumbnail(format!("attachment://{}", &attachement_name));
}
Ok((embed, icon_attachement))
}
// }}}
@ -491,7 +523,7 @@ pub struct ImageCropper {
}
impl ImageCropper {
fn crop_image_to_bytes(
pub fn crop_image_to_bytes(
&mut self,
image: &DynamicImage,
rect: AbsoluteRect,
@ -512,15 +544,48 @@ impl ImageCropper {
.to_absolute(),
)?;
let mut results = vec![];
for mode in [
PageSegMode::PsmSingleWord,
PageSegMode::PsmRawLine,
PageSegMode::PsmSingleLine,
] {
let result = self.read_score_with_mode(image, mode)?;
results.push(result.0);
// OCR sometimes loses digits
if result.0 < 1_000_000 {
continue;
} else {
return Ok(result);
}
}
Err(format!(
"Cannot read score, no matter the mode. Attempts: {:?}",
results
))?;
unreachable!()
}
pub fn read_score_with_mode(
&mut self,
image: &DynamicImage,
mode: PageSegMode,
) -> 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_image_from_mem(&self.bytes)?;
t.set_page_seg_mode(PageSegMode::PsmRawLine);
t.set_page_seg_mode(mode);
t = t.recognize()?;
let conf = t.mean_text_conf();
if t.mean_text_conf() < 10 {
Err("Score text is not readable.")?;
if conf < 10 && conf != 0 {
Err(format!(
"Score text is not readable (confidence = {}, text = {}).",
conf,
t.get_text()?.trim()
))?;
}
let text: String = t
@ -566,35 +631,103 @@ impl ImageCropper {
Ok(difficulty)
}
// }}}
// {{{ Read song
pub fn read_song(
&mut self,
image: &DynamicImage,
cache: &Mutex<SongCache>,
) -> Result<CachedSong, Error> {
self.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), title_rects())
.ok_or_else(|| "Could not find title area in picture")?
.to_absolute(),
)?;
let mut t = Tesseract::new(None, Some("eng"))?
.set_variable(
"tessedit_char_whitelist",
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ",
)?
.set_image_from_mem(&self.bytes)?;
t.set_page_seg_mode(PageSegMode::PsmSingleLine);
t = t.recognize()?;
// if t.mean_text_conf() < 10 {
// Err("Difficulty text is not readable.")?;
// }
let raw_text: &str = &t.get_text()?;
let raw_text = raw_text.trim(); // not quite raw 🤔
let mut text = raw_text;
println!("Raw text: {}, confidence: {}", text, t.mean_text_conf());
let lock = cache.lock().map_err(|_| "Poisoned song cache")?;
let cached_song = loop {
let (closest, distance) = lock
.songs()
.map(|item| {
(
item,
edit_distance::edit_distance(
&item.song.title.to_lowercase(),
&text.to_lowercase(),
),
)
})
.min_by_key(|(_, d)| *d)
.ok_or_else(|| "Empty song cache")?;
if distance > closest.song.title.len() / 3 {
if text.len() == 1 {
Err(format!(
"Could not find match for chart name '{}'",
raw_text
))?;
} else {
text = &text[..text.len() - 1];
}
} else {
break closest;
};
};
// NOTE: this will reallocate a few strings, but it is what it is
Ok(cached_song.clone())
}
// }}}
// {{{ Read jacket
pub fn read_jacket<'a>(
&mut self,
ctx: &'a UserContext,
ctx: &UserContext,
image: &DynamicImage,
) -> Result<(&'a Jacket, &'a CachedSong), Error> {
) -> Result<CachedSong, Error> {
let rect =
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), jacket_rects())
.ok_or_else(|| "Could not find jacket area in picture")?
.to_absolute();
let cropped = image.view(rect.x, rect.y, rect.width, rect.height);
let (distance, jacket) = ctx
let (distance, song_id) = ctx
.jacket_cache
.recognise(&*cropped)
.ok_or_else(|| "Could not recognise jacket")?;
if distance > 100.0 {
// Save image to be sent to discord
self.crop_image_to_bytes(&image, rect)?;
Err("No known jacket looks like this")?;
}
let song = ctx
.song_cache
.lookup(jacket.song_id)
.ok_or_else(|| format!("Could not find song with id {}", jacket.song_id))?;
.lock()
.map_err(|_| "Poisoned song cache")?
.lookup(*song_id)
.ok_or_else(|| format!("Could not find song with id {}", song_id))?
// NOTE: this will reallocate a few strings, but it is what it is
.clone();
Ok((jacket, song))
Ok(song)
}
// }}}
}