First attempt at reading the note distribution
Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
parent
49d50bf88b
commit
8339ce7054
2
.gitignore
vendored
2
.gitignore
vendored
|
@ -3,5 +3,7 @@ target
|
|||
.envrc
|
||||
data/db.sqlite
|
||||
data/jackets
|
||||
data/songs
|
||||
backups
|
||||
dump.sql
|
||||
logs
|
||||
|
|
1844
data/charts.csv
1844
data/charts.csv
File diff suppressed because it is too large
Load diff
|
@ -1,3 +0,0 @@
|
|||
Grievous Lady,grievous-lady,
|
||||
Einherjar Joker,einherjar-joker,
|
||||
Einherjar Joker,einherjar-joker-byd,BYD
|
|
|
@ -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
|
||||
|
|
|
14
schema.sql
14
schema.sql
|
@ -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);
|
||||
|
|
140
scripts/main.py
140
scripts/main.py
|
@ -16,80 +16,93 @@ 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(
|
||||
"""
|
||||
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(
|
||||
song_id = conn.execute(
|
||||
"""
|
||||
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant, jacket)
|
||||
INSERT INTO songs(title,artist,pack,side,bpm)
|
||||
VALUES (?,?,?,?,?)
|
||||
RETURNING id
|
||||
""",
|
||||
(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(" ")
|
||||
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO charts(song_id, difficulty, level, note_count, chart_constant, note_design)
|
||||
VALUES(?,?,?,?,?, ?)
|
||||
""",
|
||||
(
|
||||
song_id,
|
||||
difficulty,
|
||||
level,
|
||||
int(note_count.replace(",", "").replace(".", "")),
|
||||
int(float(cc) * 100),
|
||||
jacket,
|
||||
),
|
||||
(
|
||||
song_id,
|
||||
difficulty,
|
||||
level,
|
||||
int(note_count.replace(",", "").replace(".", "")),
|
||||
int(float(cc) * 100),
|
||||
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
30
scripts/prepare-songs.sh
Executable 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
|
69
src/chart.rs
69
src/chart.rs
|
@ -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();
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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(())
|
||||
}
|
||||
//
|
||||
// }}}
|
||||
|
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
|
|
108
src/jacket.rs
108
src/jacket.rs
|
@ -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,21 +81,99 @@ 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()?;
|
||||
entries.push((ImageVec::from_image(&image), song_id))
|
||||
match image::io::Reader::open(path) {
|
||||
Ok(reader) => {
|
||||
let image = reader.decode()?;
|
||||
entries.push((ImageVec::from_image(&image), song_id))
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let result = Self {
|
||||
|
|
251
src/score.rs
251
src/score.rs
|
@ -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,44 +855,118 @@ 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];
|
||||
|
@ -887,15 +974,20 @@ pub async fn guess_chart_name(
|
|||
} else if close_enough.len() == 1 {
|
||||
break (close_enough[0].0, close_enough[0].1);
|
||||
} else {
|
||||
Err(format!(
|
||||
"Name '{}' is too vague to choose a match",
|
||||
raw_text
|
||||
))?;
|
||||
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",
|
||||
raw_text
|
||||
))?;
|
||||
};
|
||||
};
|
||||
};
|
||||
|
||||
// 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))
|
||||
}
|
||||
// }}}
|
||||
}
|
||||
|
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue