1
Fork 0

First attempt at reading the note distribution

Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
prescientmoon 2024-07-01 18:00:03 +02:00
parent 49d50bf88b
commit 8339ce7054
Signed by: prescientmoon
SSH key fingerprint: SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
14 changed files with 1182 additions and 1488 deletions

2
.gitignore vendored
View file

@ -3,5 +3,7 @@ target
.envrc
data/db.sqlite
data/jackets
data/songs
backups
dump.sql
logs

File diff suppressed because it is too large Load diff

View file

@ -1,3 +0,0 @@
Grievous Lady,grievous-lady,
Einherjar Joker,einherjar-joker,
Einherjar Joker,einherjar-joker-byd,BYD
1 Grievous Lady grievous-lady
2 Einherjar Joker einherjar-joker
3 Einherjar Joker einherjar-joker-byd BYD

View file

@ -1 +1,49 @@
Misdeed -la bonté de Dieu et l'origine du mal-, Misdeed -la bonte de Dieu et lorigine du mal-
Name,Difficulty,Artist,Shorthand
Kanjou no Matenrou Arr.Demetori,,,matenrou
nέo κósmo,,,neo kosmo
Genesis,,Morrigan feat. Lily,genesischunithm
The Survivor (Game Edit),,,thesurvivor
Haze of Autumn,,,akinokagerou
Bamboo,,,take
10pt8ion,,,tempationgc
HIVEMIND INTERLINKED,,,hivemindrmx
Kanbu de Tomatte Sugu Tokeru,,,overdrive
1F√,,,onefr
[X],,,infinity
0xe0e1ccull,,,ifirmx
Last | Moment,,,last
Last | Eternity,,,lasteternity
Lost Emotion feat. nomico,,,lostemotion
µ,,,mu
I've heard it said,,,hearditsaid
Quon,,DJ Noriken,quonwacca
B.B.K.K.B.K.K.,,,bbkkbkk
World Fragments III(radio edit),,,worldfragments
#1f1e33,,,ifi
On And On!! feat. Jenga,,,onandon
Misdeed -la bonté de Dieu et l'origine du mal-,,,gou
Hidden Rainbows of Epicurus,,,epicurus
G e n g a o z o,,,gengaozo
A Wandering Melody of Love,,,melodyoflove
Let's Rock (Arcaea mix),,,letsrock
LunarOrbit -believe in the Espebranch road-,,,espebranch
Can I Friend You on Bassbook? Lol,,,bassline
Sheriruth (Laur Remix),,,sheriruthrm
ω4,,,omegafour
〇、,,,ichirin
Let you DIVE! (nitro rmx),,,letyoudivermx
99 Glooms,,,nnglooms
͟͝͞Ⅱ́̕,,,ii
Redraw the Colorless World,,,mukinshitsu
Ävril -Flicka i krans-,,,avril
7thSense,,,seventhsense
LIVHT MY WΔY,,,lightmyway
False Embellishment,,,kyogenkigo
Illegal Paradise,,,darakunosono
cry of viyella,,,viyella
"Good bye, Merry-Go-Round.",,,goodbyemerry
ΟΔΥΣΣΕΙΑ,,,odysseia
Mistempered Malignance,,,mismal
Twilight Concerto,,,tasogare
Heart,,,kokoro
Dancin' on a Cat's Paw,,,nekonote

1 Misdeed -la bonté de Dieu et l'origine du mal- Name Misdeed -la bonte de Dieu et lorigine du mal- Difficulty Artist Shorthand
2 Kanjou no Matenrou ~Arr.Demetori matenrou
3 nέo κósmo neo kosmo
4 Genesis Morrigan feat. Lily genesischunithm
5 The Survivor (Game Edit) thesurvivor
6 Haze of Autumn akinokagerou
7 Bamboo take
8 10pt8ion tempationgc
9 HIVEMIND INTERLINKED hivemindrmx
10 Kanbu de Tomatte Sugu Tokeru overdrive
11 1F√ onefr
12 [X] infinity
13 0xe0e1ccull ifirmx
14 Last | Moment last
15 Last | Eternity lasteternity
16 Lost Emotion feat. nomico lostemotion
17 µ mu
18 I've heard it said hearditsaid
19 Quon DJ Noriken quonwacca
20 B.B.K.K.B.K.K. bbkkbkk
21 World Fragments III(radio edit) worldfragments
22 #1f1e33 ifi
23 On And On!! feat. Jenga onandon
24 Misdeed -la bonté de Dieu et l'origine du mal- gou
25 Hidden Rainbows of Epicurus epicurus
26 G e n g a o z o gengaozo
27 A Wandering Melody of Love melodyoflove
28 Let's Rock (Arcaea mix) letsrock
29 LunarOrbit -believe in the Espebranch road- espebranch
30 Can I Friend You on Bassbook? Lol bassline
31 Sheriruth (Laur Remix) sheriruthrm
32 ω4 omegafour
33 〇、 ichirin
34 Let you DIVE! (nitro rmx) letyoudivermx
35 99 Glooms nnglooms
36 ͟͝͞Ⅱ́̕ ii
37 Redraw the Colorless World mukinshitsu
38 Ävril -Flicka i krans- avril
39 7thSense seventhsense
40 LIVHT MY WΔY lightmyway
41 False Embellishment kyogenkigo
42 Illegal Paradise darakunosono
43 cry of viyella viyella
44 Good bye, Merry-Go-Round. goodbyemerry
45 ΟΔΥΣΣΕΙΑ odysseia
46 Mistempered Malignance mismal
47 Twilight Concerto tasogare
48 Heart kokoro
49 Dancin' on a Cat's Paw nekonote

View file

@ -1,16 +1,17 @@
# {{{ users
create table IF NOT EXISTS users (
id INTEGER NOT NULL PRIMARY KEY,
discord_id TEXT UNIQUE NOT NULL,
nickname TEXT UNIQUE
discord_id TEXT UNIQUE NOT NULL
);
# }}}
# {{{ songs
CREATE TABLE IF NOT EXISTS songs (
id INTEGER NOT NULL PRIMARY KEY,
title TEXT NOT NULL,
ocr_alias TEXT,
artist TEXT,
artist TEXT NOT NULL,
side TEXT NOT NULL CHECK (side IN ('light', 'conflict', 'silent')),
bpm TEXT NOT NULL,
pack TEXT,
UNIQUE(title, artist)
);
@ -19,7 +20,8 @@ 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,
note_design TEXT,
shorthand TEXT,
difficulty TEXT NOT NULL CHECK (difficulty IN ('PST','PRS','FTR','ETR','BYD')),
level TEXT NOT NULL,
@ -53,4 +55,4 @@ CREATE TABLE IF NOT EXISTS plays (
);
# }}}
insert into users(discord_id, nickname) values (385759924917108740, 'prescientmoon');
insert into users(discord_id) values (385759924917108740);

View file

@ -16,65 +16,49 @@ conn = sqlite3.connect(db_path)
# {{{ Import songs
def import_charts_from_csv():
chart_count = 0
songs = dict()
song_count = 0
shorthand_count = 0
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] = {"charts": [], "shorthand": None}
songs[title]["charts"].append([difficulty, level, cc, note_count, None])
for i, row in enumerate(csv.reader(file)):
if i == 0 or len(row) == 0:
continue
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
song_count += 1
[
title,
artist,
pack,
*charts,
side,
bpm,
version,
date,
ext_version,
ext_date,
original,
] = map(lambda v: v.strip().replace("\n", " "), row)
for i in range(len(songs[title]["charts"])):
if songs[title]["charts"][i][0] == difficulty:
songs[title]["charts"][i][4] = jacket
changed += 1
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
# Problematic titles that can belong to multiple artists
for possibility in ["Quon", "Gensis"]:
if title.startswith(possibility):
artist = title[len(possibility) + 2 : -1]
title = possibility
break
row = conn.execute(
song_id = conn.execute(
"""
INSERT INTO songs(title,artist,ocr_alias)
VALUES (?,?,?)
INSERT INTO songs(title,artist,pack,side,bpm)
VALUES (?,?,?,?,?)
RETURNING id
""",
(title, artist, entry.get("shorthand")),
).fetchone()
song_id = row[0]
(title, artist, pack, side.lower(), bpm),
).fetchone()[0]
for i in range(4):
[note_design, level, cc, note_count] = charts[i * 4 : (i + 1) * 4]
if note_design == "N/A":
continue
chart_count += 2
[difficulty, level] = level.split(" ")
for difficulty, level, cc, note_count, jacket in entry["charts"]:
conn.execute(
"""
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant, jacket)
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant, note_design)
VALUES(?,?,?,?,?, ?)
""",
(
@ -83,13 +67,42 @@ def import_charts_from_csv():
level,
int(note_count.replace(",", "").replace(".", "")),
int(float(cc) * 100),
jacket,
note_design if len(note_design) else None,
),
)
with open(data_dir + "/shorthands.csv", mode="r") as file:
for i, row in enumerate(csv.reader(file)):
if i == 0 or len(row) == 0:
continue
shorthand_count += 1
[name, difficulty, artist, shorthand] = map(lambda v: v.strip(), row)
conn.execute(
f"""
UPDATE charts
SET shorthand=?
WHERE EXISTS (
SELECT 1 FROM songs s
WHERE s.id = charts.song_id
AND s.title=?
{"" if artist=="" else "AND artist=?"}
)
{"" if difficulty=="" else "AND difficulty=?"}
""",
[
shorthand,
name,
*([] if artist == "" else [artist]),
*([] if difficulty == "" else [difficulty]),
],
)
conn.commit()
print(f"Imported {chart_count} charts and {len(songs)} songs")
print(
f"Imported {chart_count} charts, {song_count} songs, and {shorthand_count} shorthands"
)
# }}}
@ -99,6 +112,5 @@ subcommand = sys.argv[2]
if command == "import" and subcommand == "charts":
import_charts_from_csv()
&song_title
if command == "export" and subcommand == "jackets":
import_charts_from_csv()

30
scripts/prepare-songs.sh Executable file
View file

@ -0,0 +1,30 @@
#!/usr/bin/env bash
dir_path=./data/songs
# Find all files in the directory and its subdirectories
find "$dir_path" -type f | while read -r file; do
# Get the filename without the directory path
filename=$(basename "$file")
# Check if the filename starts with "1080_"
if [[ $filename == 1080_* ]]; then
# Remove the "1080_" prefix
new_filename="${filename#1080_}"
# Get the directory path without the filename
file_dir=$(dirname "$file")
# Construct the new file path
new_file_path="$file_dir/$new_filename"
# Rename the file
mv "$file" "$new_file_path"
echo "Renamed: $file -> $new_file_path"
fi
done
mv $dir_path/dropdead/3*.jpg $dir_path/overdead 2>/dev/null
mv $dir_path/singularity/3*.jpg $dir_path/singularityvvvip 2>/dev/null
mv $dir_path/redandblue/3*.jpg $dir_path/redandblueandgreen 2>/dev/null
mv $dir_path/ignotus/3*.jpg $dir_path/ignotusafterburn 2>/dev/null
rm -rf $dir_path/ifirmxrmx 2>/dev/null

View file

@ -18,7 +18,9 @@ impl Difficulty {
pub const DIFFICULTIES: [Difficulty; 5] =
[Self::PST, Self::PRS, Self::FTR, Self::ETR, Self::BYD];
pub const DIFFICULTY_STRINGS: [&'static str; 5] = ["PST", "PRS", "FTR", "ETR", "BYD"];
pub const DIFFICULTY_SHORTHANDS: [&'static str; 5] = ["PST", "PRS", "FTR", "ETR", "BYD"];
pub const DIFFICULTY_STRINGS: [&'static str; 5] =
["past", "present", "future", "eternal", "beyond"];
#[inline]
pub fn to_index(self) -> usize {
@ -30,7 +32,7 @@ impl TryFrom<String> for Difficulty {
type Error = String;
fn try_from(value: String) -> Result<Self, Self::Error> {
for (i, s) in Self::DIFFICULTY_STRINGS.iter().enumerate() {
for (i, s) in Self::DIFFICULTY_SHORTHANDS.iter().enumerate() {
if value == **s {
return Ok(Self::DIFFICULTIES[i]);
}
@ -45,15 +47,8 @@ impl TryFrom<String> for Difficulty {
pub struct Song {
pub id: u32,
pub title: String,
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)
}
#[allow(dead_code)]
pub artist: String,
}
// }}}
// {{{ Chart
@ -61,6 +56,7 @@ impl Song {
pub struct Chart {
pub id: u32,
pub song_id: u32,
pub shorthand: Option<String>,
pub difficulty: Difficulty,
pub level: String, // TODO: this could become an enum
@ -68,7 +64,16 @@ pub struct Chart {
pub note_count: u32,
pub chart_constant: u32,
pub jacket: Option<PathBuf>,
pub cached_jacket: Option<&'static [u8]>,
}
impl Chart {
#[inline]
pub fn jacket_path(&self, data_dir: &PathBuf) -> PathBuf {
data_dir
.join("jackets")
.join(format!("{}-{}.jpg", self.song_id, self.id))
}
}
// }}}
// {{{ Cached song
@ -116,6 +121,11 @@ impl CachedSong {
pub fn charts(&self) -> impl Iterator<Item = &Chart> {
self.charts.iter().filter_map(|i| i.as_ref())
}
#[inline]
pub fn charts_mut(&mut self) -> impl Iterator<Item = &mut Chart> {
self.charts.iter_mut().filter_map(|i| i.as_mut())
}
}
// }}}
// {{{ Song cache
@ -126,8 +136,11 @@ pub struct SongCache {
impl SongCache {
#[inline]
pub fn lookup(&self, id: u32) -> Option<&CachedSong> {
self.songs.get(id as usize).and_then(|i| i.as_ref())
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]
@ -153,13 +166,33 @@ impl SongCache {
.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(data_dir: &PathBuf, pool: &SqlitePool) -> Result<Self, Error> {
pub async fn new(pool: &SqlitePool) -> Result<Self, Error> {
let mut result = Self::default();
let songs = sqlx::query!("SELECT * FROM songs").fetch_all(pool).await?;
@ -168,7 +201,6 @@ impl SongCache {
let song = Song {
id: song.id as u32,
title: song.title,
ocr_alias: song.ocr_alias,
artist: song.artist,
};
@ -187,13 +219,12 @@ impl SongCache {
let chart = Chart {
id: chart.id as u32,
song_id: chart.song_id as u32,
shorthand: chart.shorthand,
difficulty: Difficulty::try_from(chart.difficulty)?,
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))),
cached_jacket: None,
};
let index = chart.difficulty.to_index();

View file

@ -6,11 +6,9 @@ use crate::score::{
};
use crate::user::{discord_it_to_discord_user, User};
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
/// Show this help menu
@ -118,8 +116,6 @@ pub async fn magic(
.content(format!("Image {}: reading jacket", i + 1));
handle.edit(ctx, edited).await?;
let song_by_jacket = cropper.read_jacket(ctx.data(), &image).await;
// This makes OCR more likely to work
let mut ocr_image = image.grayscale().blur(1.);
@ -146,6 +142,9 @@ pub async fn magic(
Ok(d) => d,
};
let song_by_jacket = cropper.read_jacket(ctx.data(), &image, difficulty).await;
let note_distribution = cropper.read_distribution(&image)?;
ocr_image.invert();
let edited = CreateReply::default()
@ -153,12 +152,20 @@ pub async fn magic(
.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, difficulty)
.await;
let song_by_name =
cropper.read_song(&ocr_image, &ctx.data().song_cache, difficulty);
let (song, chart) = match (song_by_jacket, song_by_name) {
// {{{ Both errors
(Err(err_jacket), Err(err_name)) => {
cropper.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(
ImageDimensions::from_image(&image),
jacket_rects(),
)
.ok_or_else(|| "Could not find jacket area in picture")?
.to_absolute(),
)?;
error_with_image(
ctx,
&cropper.bytes,
@ -189,60 +196,8 @@ Title error: {}
}
// }}}
// {{{ Only name succeeded
(Err(err_jacket), Ok(mut by_name)) => {
(Err(err_jacket), Ok(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();
// }}}
// {{{ Build path
let filename = format!("{}-{}", by_name.0.id, by_name.1.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,
by_name.1.song_id,
by_name.1.difficulty,
)
.execute(&ctx.data().db)
.await?;
// }}}
// {{{ Aquire and use song cache lock
{
let mut song_cache = ctx.data().song_cache.lock().await;
let chart = song_cache
.lookup_mut(by_name.0.id)?
.lookup_mut(difficulty)?;
if chart.jacket.is_none() {
by_name.1.jacket = Some(jacket_path.clone());
chart.jacket = Some(jacket_path);
} else {
println!(
"Jacket not detected for chart {} [{:?}]",
by_name.0.id, difficulty
)
};
}
// }}}
by_name
}
// }}}
@ -250,8 +205,8 @@ Title error: {}
(Ok(by_jacket), Ok(by_name)) => {
if by_name.0.id != by_jacket.0.id {
println!(
"Got diverging choices between '{:?}' and '{:?}'",
by_jacket.0.id, by_name.0.id
"Got diverging choices between '{}' and '{}'",
by_jacket.0.title, by_name.0.title
);
};
@ -284,8 +239,21 @@ Title error: {}
};
// {{{ Build play
let (score, maybe_fars, score_warning) =
Score::resolve_ambiguities(score_possibilities, None, chart.note_count)?;
let (score, maybe_fars, score_warning) = Score::resolve_ambiguities(
score_possibilities,
Some(note_distribution),
chart.note_count,
)
.map_err(|err| {
format!(
"Error occurred when disambiguating scores for '{}' [{:?}] by {}: {}",
song.title, difficulty, song.artist, err
)
})?;
println!(
"Maybe fars {:?}, distribution {:?}",
maybe_fars, note_distribution
);
let play = CreatePlay::new(score, &chart, &user)
.with_attachment(file)
.with_fars(maybe_fars)
@ -382,8 +350,6 @@ pub async fn show(
return Ok(());
}
let lock = ctx.data().song_cache.lock().await;
let mut embeds = Vec::with_capacity(ids.len());
let mut attachments = Vec::with_capacity(ids.len());
for (i, id) in ids.iter().enumerate() {
@ -419,7 +385,7 @@ pub async fn show(
let user = discord_it_to_discord_user(&ctx, &res.discord_id).await?;
let (song, chart) = lock.lookup_chart(play.chart_id)?;
let (song, chart) = ctx.data().song_cache.lookup_chart(play.chart_id)?;
let (embed, attachment) = play.to_embed(song, chart, i, Some(&user)).await?;
embeds.push(embed);

View file

@ -1,6 +1,6 @@
use std::io::Cursor;
use chrono::{DateTime, NaiveDateTime};
use chrono::DateTime;
use image::{ImageBuffer, Rgb};
use plotters::{
backend::{BitMapBackend, PixelFormat, RGBPixel},
@ -8,10 +8,7 @@ use plotters::{
drawing::IntoDrawingArea,
element::Circle,
series::LineSeries,
style::{
text_anchor::{HPos, Pos, VPos},
Color, FontTransform, IntoFont, TextStyle, BLUE, WHITE,
},
style::{Color, IntoFont, TextStyle, BLUE, WHITE},
};
use poise::{
serenity_prelude::{CreateAttachment, CreateMessage},
@ -72,13 +69,18 @@ pub async fn best(
let (name, difficulty) = name
.strip_suffix("PST")
.zip(Some(Difficulty::PST))
.or_else(|| name.strip_suffix("[PST]").zip(Some(Difficulty::PST)))
.or_else(|| name.strip_suffix("PRS").zip(Some(Difficulty::PRS)))
.or_else(|| name.strip_suffix("[PRS]").zip(Some(Difficulty::PRS)))
.or_else(|| name.strip_suffix("FTR").zip(Some(Difficulty::FTR)))
.or_else(|| name.strip_suffix("[FTR]").zip(Some(Difficulty::FTR)))
.or_else(|| name.strip_suffix("ETR").zip(Some(Difficulty::ETR)))
.or_else(|| name.strip_suffix("[ETR]").zip(Some(Difficulty::ETR)))
.or_else(|| name.strip_suffix("BYD").zip(Some(Difficulty::BYD)))
.or_else(|| name.strip_suffix("[BYD]").zip(Some(Difficulty::BYD)))
.unwrap_or((&name, Difficulty::FTR));
let (song, chart) = guess_chart_name(name, &ctx.data().song_cache, difficulty).await?;
let (song, chart) = guess_chart_name(name, &ctx.data().song_cache, Some(difficulty), true)?;
let play = query_as!(
DbPlay,
@ -93,7 +95,12 @@ pub async fn best(
)
.fetch_one(&ctx.data().db)
.await
.map_err(|_| format!("Could not find any scores for chart"))?
.map_err(|_| {
format!(
"Could not find any scores for {} [{:?}]",
song.title, chart.difficulty
)
})?
.to_play();
let (embed, attachment) = play
@ -112,7 +119,7 @@ pub async fn best(
Ok(())
}
// }}}
// Score plot
// {{{ Score plot
/// Show the best score on a given chart
#[poise::command(prefix_command, slash_command)]
pub async fn plot(
@ -134,13 +141,18 @@ pub async fn plot(
let (name, difficulty) = name
.strip_suffix("PST")
.zip(Some(Difficulty::PST))
.or_else(|| name.strip_suffix("[PST]").zip(Some(Difficulty::PST)))
.or_else(|| name.strip_suffix("PRS").zip(Some(Difficulty::PRS)))
.or_else(|| name.strip_suffix("[PRS]").zip(Some(Difficulty::PRS)))
.or_else(|| name.strip_suffix("FTR").zip(Some(Difficulty::FTR)))
.or_else(|| name.strip_suffix("[FTR]").zip(Some(Difficulty::FTR)))
.or_else(|| name.strip_suffix("ETR").zip(Some(Difficulty::ETR)))
.or_else(|| name.strip_suffix("[ETR]").zip(Some(Difficulty::ETR)))
.or_else(|| name.strip_suffix("BYD").zip(Some(Difficulty::BYD)))
.or_else(|| name.strip_suffix("[BYD]").zip(Some(Difficulty::BYD)))
.unwrap_or((&name, Difficulty::FTR));
let (song, chart) = guess_chart_name(name, &ctx.data().song_cache, difficulty).await?;
let (song, chart) = guess_chart_name(name, &ctx.data().song_cache, Some(difficulty), true)?;
let plays = query_as!(
DbPlay,
@ -157,7 +169,11 @@ pub async fn plot(
.await?;
if plays.len() == 0 {
ctx.reply("No plays found").await?;
ctx.reply(format!(
"No plays found on {} [{:?}]",
song.title, chart.difficulty
))
.await?;
return Ok(());
}
@ -182,7 +198,7 @@ pub async fn plot(
let mut buffer = vec![u8::MAX; RGBPixel::PIXEL_SIZE * (width * height) as usize];
{
let mut root = BitMapBackend::with_buffer(&mut buffer, (width, height)).into_drawing_area();
let root = BitMapBackend::with_buffer(&mut buffer, (width, height)).into_drawing_area();
let mut chart = ChartBuilder::on(&root)
.margin(25)
@ -242,4 +258,4 @@ pub async fn plot(
Ok(())
}
//
// }}}

View file

@ -1,7 +1,6 @@
use std::{path::PathBuf, sync::Arc};
use std::path::PathBuf;
use sqlx::SqlitePool;
use tokio::sync::Mutex;
use crate::{chart::SongCache, jacket::JacketCache};
@ -11,21 +10,25 @@ pub type Context<'a> = poise::Context<'a, UserContext, Error>;
// Custom user data passed to all command functions
pub struct UserContext {
#[allow(dead_code)]
pub data_dir: PathBuf,
pub db: SqlitePool,
pub song_cache: Arc<Mutex<SongCache>>,
pub song_cache: 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(&data_dir, &db).await?;
let jacket_cache = JacketCache::new(&song_cache)?;
let mut song_cache = SongCache::new(&db).await?;
let jacket_cache = JacketCache::new(&data_dir, &mut song_cache)?;
println!("Created user context");
Ok(Self {
data_dir,
db,
song_cache: Arc::new(Mutex::new(song_cache)),
song_cache,
jacket_cache,
})
}

View file

@ -1,14 +1,18 @@
use std::{collections::HashSet, path::PathBuf};
use std::{collections::HashSet, fs, path::PathBuf, str::FromStr};
use image::{GenericImageView, Rgba};
use kd_tree::{KdMap, KdPoint};
use num::Integer;
use crate::{chart::SongCache, context::Error};
use crate::{
chart::{Difficulty, SongCache},
context::Error,
score::guess_chart_name,
};
/// How many sub-segments to split each side into
const SPLIT_FACTOR: u32 = 5;
const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize;
pub const SPLIT_FACTOR: u32 = 8;
pub const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize;
#[derive(Debug, Clone)]
pub struct ImageVec {
@ -77,22 +81,100 @@ pub struct JacketCache {
impl JacketCache {
// {{{ Generate tree
// 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();
pub fn new(data_dir: &PathBuf, song_cache: &mut SongCache) -> Result<Self, Error> {
let jacket_dir = data_dir.join("jackets");
for item in song_cache.songs() {
for chart in item.charts() {
if let Some(jacket) = &chart.jacket {
jackets.insert((jacket, item.song.id));
if jacket_dir.exists() {
fs::remove_dir_all(&jacket_dir).expect("Could not delete jacket dir");
}
fs::create_dir_all(&jacket_dir).expect("Could not create jacket dir");
let mut jackets: HashSet<(PathBuf, u32)> = HashSet::new();
let entries = fs::read_dir(data_dir.join("songs")).expect("Couldn't read songs directory");
for entry in entries {
let dir = entry?;
let raw_dir_name = dir.file_name();
let dir_name = raw_dir_name.to_str().unwrap();
for entry in fs::read_dir(dir.path()).expect("Couldn't read song directory") {
let file = entry?;
let raw_name = file.file_name();
let name = raw_name.to_str().unwrap().strip_suffix(".jpg").unwrap();
if !name.ends_with("_256") {
continue;
}
let name = name.strip_suffix("_256").unwrap();
let difficulty = match name {
"0" => Some(Difficulty::PST),
"1" => Some(Difficulty::PRS),
"2" => Some(Difficulty::FTR),
"3" => Some(Difficulty::BYD),
"4" => Some(Difficulty::ETR),
"base" => None,
"base_night" => None,
"base_ja" => None,
_ => Err(format!("Unknown jacket suffix {}", name))?,
};
let (song, chart) = guess_chart_name(dir_name, &song_cache, difficulty, true)?;
jackets.insert((file.path(), song.id));
let contents = fs::read(file.path())?.leak();
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");
chart.cached_jacket = Some(contents);
}
}
} 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();
chart.cached_jacket = Some(contents);
}
}
}
let mut entries = vec![];
for (path, song_id) in jackets {
let image = image::io::Reader::open(path)?.decode()?;
match image::io::Reader::open(path) {
Ok(reader) => {
let image = reader.decode()?;
entries.push((ImageVec::from_image(&image), song_id))
}
_ => continue,
}
}
let result = Self {
tree: KdMap::build_by_ordered_float(entries),

View file

@ -1,5 +1,6 @@
#![allow(dead_code)]
use std::fmt::Display;
use std::fs;
use std::io::Cursor;
use std::str::FromStr;
use std::sync::OnceLock;
@ -11,10 +12,10 @@ use poise::serenity_prelude::{
Attachment, AttachmentId, CreateAttachment, CreateEmbed, CreateEmbedAuthor, Timestamp,
};
use tesseract::{PageSegMode, Tesseract};
use tokio::sync::Mutex;
use crate::chart::{Chart, Difficulty, Song, SongCache};
use crate::context::{Error, UserContext};
use crate::jacket::IMAGE_VEC_DIM;
use crate::user::User;
// {{{ Score
@ -183,7 +184,7 @@ impl Score {
};
// }}}
if scores.len() == 0 {
if scores.len() == 1 {
Ok((score, consensus_fars, None))
} else {
Ok((score, consensus_fars, Some("Due to a reading error, I could not make sure the shiny-amount I calculated is accurate!")))
@ -506,11 +507,8 @@ impl Play {
author: Option<&poise::serenity_prelude::User>,
) -> Result<(CreateEmbed, Option<CreateAttachment>), Error> {
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)
.await?,
),
let icon_attachement = match chart.cached_jacket {
Some(bytes) => Some(CreateAttachment::bytes(bytes, &attachement_name)),
None => None,
};
@ -545,7 +543,7 @@ impl Play {
true,
)
.field("Max recall", "?", true)
.field("Id", format!("{}", self.id), true);
.field("ID", format!("{}", self.id), true);
if icon_attachement.is_some() {
embed = embed.thumbnail(format!("attachment://{}", &attachement_name));
@ -679,6 +677,20 @@ impl RelativeRect {
}
}
/// Shift this rect on the y axis by a given absolute pixel amount
#[inline]
pub fn shift_y_abs(&self, amount: u32) -> Self {
let mut res = Self::new(
self.x,
self.y + (amount as f32 / self.dimensions.height as f32),
self.width,
self.height,
self.dimensions,
);
res.fix();
res
}
/// Clamps the values apropriately
#[inline]
pub fn fix(&mut self) {
@ -795,6 +807,7 @@ fn difficulty_rects() -> &'static [RelativeRect] {
AbsoluteRect::new(183, 487, 197, 44, ImageDimensions::new(2560, 1600)).to_relative(),
AbsoluteRect::new(198, 692, 219, 46, ImageDimensions::new(2732, 2048)).to_relative(),
AbsoluteRect::new(414, 364, 177, 38, ImageDimensions::new(2778, 1284)).to_relative(),
AbsoluteRect::new(76, 172, 77, 18, ImageDimensions::new(1080, 607)).to_relative(),
];
process_datapoints(&mut rects);
rects
@ -842,50 +855,128 @@ pub fn jacket_rects() -> &'static [RelativeRect] {
})
}
// }}}
// {{{ Note distribution
pub fn note_distribution_rects() -> (
&'static [RelativeRect],
&'static [RelativeRect],
&'static [RelativeRect],
) {
static CELL: OnceLock<(
&'static [RelativeRect],
&'static [RelativeRect],
&'static [RelativeRect],
)> = OnceLock::new();
*CELL.get_or_init(|| {
let mut pure_rects: Vec<RelativeRect> = vec![
AbsoluteRect::new(729, 523, 58, 22, ImageDimensions::new(1560, 720)).to_relative(),
AbsoluteRect::new(815, 520, 57, 23, ImageDimensions::new(1600, 720)).to_relative(),
AbsoluteRect::new(1019, 856, 91, 33, ImageDimensions::new(2000, 1200)).to_relative(),
AbsoluteRect::new(1100, 1085, 102, 38, ImageDimensions::new(2160, 1620)).to_relative(),
AbsoluteRect::new(1130, 1118, 105, 39, ImageDimensions::new(2224, 1668)).to_relative(),
AbsoluteRect::new(1286, 850, 91, 35, ImageDimensions::new(2532, 1170)).to_relative(),
AbsoluteRect::new(1305, 1125, 117, 44, ImageDimensions::new(2560, 1600)).to_relative(),
AbsoluteRect::new(1389, 1374, 126, 48, ImageDimensions::new(2732, 2048)).to_relative(),
AbsoluteRect::new(1407, 933, 106, 40, ImageDimensions::new(2778, 1284)).to_relative(),
];
process_datapoints(&mut pure_rects);
let skip_distances = vec![40, 40, 57, 67, 65, 60, 75, 78, 65];
let far_rects: Vec<_> = pure_rects
.iter()
.enumerate()
.map(|(i, rect)| rect.shift_y_abs(skip_distances[i]))
.collect();
let lost_rects: Vec<_> = far_rects
.iter()
.enumerate()
.map(|(i, rect)| rect.shift_y_abs(skip_distances[i]))
.collect();
(pure_rects.leak(), far_rects.leak(), lost_rects.leak())
})
}
// }}}
// }}}
// }}}
// {{{ Recognise chart name
pub async fn guess_chart_name(
/// Runs a specialized fuzzy-search through all charts in the game.
///
/// The `unsafe_heuristics` toggle increases the amount of resolvable queries, but might let in
/// some false positives. We turn it on for simple user-search commands, but disallow it for things
/// like OCR-generated text.
pub fn guess_chart_name<'a>(
raw_text: &str,
cache: &Mutex<SongCache>,
difficulty: Difficulty,
) -> Result<(Song, Chart), Error> {
cache: &'a SongCache,
difficulty: Option<Difficulty>,
unsafe_heuristics: bool,
) -> Result<(&'a Song, &'a Chart), Error> {
let raw_text = raw_text.trim(); // not quite raw 🤔
let mut text: &str = &raw_text.to_lowercase();
let lock = cache.lock().await;
// Cached vec used to store distance calculations
let mut distance_vec = Vec::with_capacity(3);
let (song, chart) = loop {
let close_enough: Vec<_> = lock
let mut close_enough: Vec<_> = cache
.songs()
.filter_map(|item| Some((&item.song, item.lookup(difficulty).ok()?)))
.map(|(song, chart)| {
let song_title = song.title.to_lowercase();
let shortest_len = Ord::min(song_title.len(), text.len());
let mut smallest_distance = edit_distance(&text, &song_title);
.filter_map(|item| {
let song = &item.song;
let chart = if let Some(difficulty) = difficulty {
item.lookup(difficulty).ok()?
} else {
item.charts().next()?
};
if let Some(sliced) = &song_title.get(..shortest_len)
&& text.len() >= 6
{
// We want to make this route super costly, which is why we multiply by 50
smallest_distance = smallest_distance.min(50 * edit_distance(&text, sliced));
let song_title = song.title.to_lowercase();
distance_vec.clear();
let base_distance = edit_distance(&text, &song_title);
if base_distance < 1.max(song.title.len() / 3) {
distance_vec.push(base_distance * 10 + 2);
}
(song, chart, smallest_distance)
let shortest_len = Ord::min(song_title.len(), text.len());
if let Some(sliced) = &song_title.get(..shortest_len)
&& (text.len() >= 6 || unsafe_heuristics)
{
let slice_distance = edit_distance(&text, sliced);
if slice_distance < 1 {
distance_vec.push(slice_distance * 10 + 3);
}
}
if let Some(shorthand) = &chart.shorthand
&& unsafe_heuristics
{
let short_distance = edit_distance(&text, shorthand);
if short_distance < 1.max(shorthand.len() / 3) {
distance_vec.push(short_distance * 10 + 1);
}
}
distance_vec
.iter()
.min()
.map(|distance| (song, chart, *distance))
})
.filter(|(song, _, d)| *d < song.title.len() / 3)
.collect();
if close_enough.len() == 0 {
if text.len() == 1 {
if text.len() <= 1 {
Err(format!(
"Could not find match for chart name '{}'",
raw_text
"Could not find match for chart name '{}' [{:?}]",
raw_text, difficulty
))?;
} else {
text = &text[..text.len() - 1];
}
} else if close_enough.len() == 1 {
break (close_enough[0].0, close_enough[0].1);
} else {
if unsafe_heuristics {
close_enough.sort_by_key(|(_, _, distance)| *distance);
break (close_enough[0].0, close_enough[0].1);
} else {
Err(format!(
"Name '{}' is too vague to choose a match",
@ -893,9 +984,10 @@ pub async fn guess_chart_name(
))?;
};
};
};
// NOTE: this will reallocate a few strings, but it is what it is
Ok((song.clone(), chart.clone()))
Ok((song, chart))
}
// }}}
// {{{ Run OCR
@ -916,6 +1008,9 @@ impl ImageCropper {
let image = image.crop_imm(rect.x, rect.y, rect.width, rect.height);
let mut cursor = Cursor::new(&mut self.bytes);
image.write_to(&mut cursor, image::ImageFormat::Png)?;
fs::write(format!("./logs/{}.png", Timestamp::now()), &self.bytes)?;
Ok(())
}
@ -956,6 +1051,7 @@ impl ImageCropper {
// so we try to detect that and fix it
loop {
let old_stack_len = results.len();
println!("Results {:?}", results);
results = results
.iter()
.flat_map(|result| {
@ -1000,6 +1096,7 @@ impl ImageCropper {
})
.map(|r| Score(r))
.collect();
println!("Results {:?}", results);
// 2. Look for consensus
for result in results.iter() {
@ -1012,6 +1109,7 @@ impl ImageCropper {
// If there's no consensus, we return everything
results.sort();
results.dedup();
println!("Results {:?}", results);
Ok(results)
}
@ -1060,7 +1158,7 @@ impl ImageCropper {
t = t.recognize()?;
let text: &str = &t.get_text()?;
let text = text.trim();
let text = text.trim().to_lowercase();
let conf = t.mean_text_conf();
if conf < 10 && conf != 0 {
@ -1073,7 +1171,7 @@ impl ImageCropper {
let difficulty = Difficulty::DIFFICULTIES
.iter()
.zip(Difficulty::DIFFICULTY_STRINGS)
.min_by_key(|(_, difficulty_string)| edit_distance(difficulty_string, text))
.min_by_key(|(_, difficulty_string)| edit_distance(difficulty_string, &text))
.map(|(difficulty, _)| *difficulty)
.ok_or_else(|| format!("Unrecognised difficulty '{}'", text))?;
@ -1081,12 +1179,12 @@ impl ImageCropper {
}
// }}}
// {{{ Read song
pub async fn read_song(
pub fn read_song<'a>(
&mut self,
image: &DynamicImage,
cache: &Mutex<SongCache>,
cache: &'a SongCache,
difficulty: Difficulty,
) -> Result<(Song, Chart), Error> {
) -> Result<(&'a Song, &'a Chart), Error> {
self.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), title_rects())
@ -1105,24 +1203,25 @@ impl ImageCropper {
let raw_text: &str = &t.get_text()?;
let conf = t.mean_text_conf();
if conf < 20 && conf != 0 {
Err(format!(
"Title text is not readable (confidence = {}, text = {}).",
conf,
raw_text.trim()
))?;
}
// let conf = t.mean_text_conf();
// if conf < 20 && conf != 0 {
// Err(format!(
// "Title text is not readable (confidence = {}, text = {}).",
// conf,
// raw_text.trim()
// ))?;
// }
guess_chart_name(raw_text, cache, difficulty).await
guess_chart_name(raw_text, cache, Some(difficulty), false)
}
// }}}
// {{{ Read jacket
pub async fn read_jacket<'a>(
&mut self,
ctx: &UserContext,
ctx: &'a UserContext,
image: &DynamicImage,
) -> Result<(Song, Chart), Error> {
difficulty: Difficulty,
) -> Result<(&'a Song, &'a Chart), Error> {
let rect =
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), jacket_rects())
.ok_or_else(|| "Could not find jacket area in picture")?
@ -1134,15 +1233,59 @@ impl ImageCropper {
.recognise(&*cropped)
.ok_or_else(|| "Could not recognise jacket")?;
if distance > 100.0 {
if distance > (IMAGE_VEC_DIM * 3) as f32 {
Err("No known jacket looks like this")?;
}
let lock = ctx.song_cache.lock().await;
let (song, chart) = lock.lookup_chart(*song_id)?;
let item = ctx.song_cache.lookup(*song_id)?;
let chart = item.lookup(difficulty)?;
// NOTE: this will reallocate a few strings, but it is what it is
Ok((song.clone(), chart.clone()))
Ok((&item.song, chart))
}
// }}}
// {{{ Read distribution
pub fn read_distribution(&mut self, image: &DynamicImage) -> Result<(u32, u32, u32), Error> {
let mut t = Tesseract::new(None, Some("eng"))?
.set_variable("classify_bln_numeric_mode", "1")?
.set_variable("tessedit_char_whitelist", "0123456789")?;
t.set_page_seg_mode(PageSegMode::PsmSingleLine);
let (pure_rects, far_rects, lost_rects) = note_distribution_rects();
self.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), pure_rects)
.ok_or_else(|| "Could not find pure-rect area in picture")?
.to_absolute(),
)?;
t = t.set_image_from_mem(&self.bytes)?.recognize()?;
let pure_notes = u32::from_str(&t.get_text()?.trim()).unwrap_or(0);
println!("Raw {}", t.get_text()?.trim());
self.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), far_rects)
.ok_or_else(|| "Could not find far-rect area in picture")?
.to_absolute(),
)?;
t = t.set_image_from_mem(&self.bytes)?.recognize()?;
let far_notes = u32::from_str(&t.get_text()?.trim()).unwrap_or(0);
println!("Raw {}", t.get_text()?.trim());
self.crop_image_to_bytes(
&image,
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), lost_rects)
.ok_or_else(|| "Could not find lost-rect area in picture")?
.to_absolute(),
)?;
t = t.set_image_from_mem(&self.bytes)?.recognize()?;
let lost_notes = u32::from_str(&t.get_text()?.trim()).unwrap_or(0);
println!("Raw {}", t.get_text()?.trim());
Ok((pure_notes, far_notes, lost_notes))
}
// }}}
}

View file

@ -8,7 +8,6 @@ use crate::context::{Context, Error};
pub struct User {
pub id: u32,
pub discord_id: String,
pub nickname: Option<String>,
}
impl User {
@ -21,7 +20,6 @@ impl User {
Ok(User {
id: user.id as u32,
discord_id: user.discord_id,
nickname: user.nickname,
})
}
}