So much work, my god
Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
parent
5cfeff4e14
commit
5c4bfa25c9
SSH key fingerprint:
SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
|
|
@ -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
data/shorthands.csv
Normal file
1
data/shorthands.csv
Normal file
|
|
@ -0,0 +1 @@
|
|||
Misdeed -la bonté de Dieu et l'origine du mal-, Misdeed -la bonte de Dieu et lorigine du mal-
|
||||
|
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
43
src/chart.rs
|
|
@ -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
235
src/commands.rs
|
|
@ -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(())
|
||||
}
|
||||
// }}}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
203
src/score.rs
|
|
@ -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)
|
||||
}
|
||||
// }}}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue