prescientmoon/shimmeringmoon
1
Fork 0
Code Activity

Initial commit, I guess

Browse source 
Signed-off-by: prescientmoon <git@moonythm.dev>
This commit is contained in:
prescientmoon 2024-06-22 16:40:56 +02:00
commit a0e3decd7a
Signed by: prescientmoon
SSH key fingerprint: SHA256:UUF9JT2s8Xfyv76b8ZuVL7XrmimH4o49p4b+iexbVH4
14 changed files with 4813 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.gitignore vendored Normal file
View file

@ -0,0 +1,4 @@
target
.direnv
.envrc
data

3729
Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

19
Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "shimmeringmoon"
version = "0.1.0"
edition = "2021"
[dependencies]
chrono = "0.4.38"
edit-distance = "2.1.0"
image = "0.25.1"
num = "0.4.3"
plotlib = "0.5.1"
poise = "0.6.1"
prettytable-rs = "0.10.0"
sqlx = { version = "0.7.4", features = ["sqlite", "runtime-tokio", "chrono"] }
tesseract = "0.15.1"
tokio = {version="1.38.0", features=["rt-multi-thread"]}
[profile.dev.package.sqlx-macros]
opt-level = 3

100
flake.lock Normal file
View file

@ -0,0 +1,100 @@
{
"nodes": {
"fenix": {
"inputs": {
"nixpkgs": [
"nixpkgs"
],
"rust-analyzer-src": "rust-analyzer-src"
},
"locked": {
"lastModified": 1717827974,
"narHash": "sha256-ixopuTeTouxqTxfMuzs6IaRttbT8JqRW5C9Q/57WxQw=",
"owner": "nix-community",
"repo": "fenix",
"rev": "ab655c627777ab5f9964652fe23bbb1dfbd687a8",
"type": "github"
},
"original": {
"owner": "nix-community",
"repo": "fenix",
"type": "github"
}
},
"flake-utils": {
"inputs": {
"systems": "systems"
},
"locked": {
"lastModified": 1710146030,
"narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=",
"owner": "numtide",
"repo": "flake-utils",
"rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a",
"type": "github"
},
"original": {
"owner": "numtide",
"repo": "flake-utils",
"type": "github"
}
},
"nixpkgs": {
"locked": {
"lastModified": 1718000748,
"narHash": "sha256-zliqz7ovpxYdKIK+GlWJZxifXsT9A1CHNQhLxV0G1Hc=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "869cab745a802b693b45d193b460c9184da671f3",
"type": "github"
},
"original": {
"owner": "nixos",
"ref": "release-24.05",
"repo": "nixpkgs",
"type": "github"
}
},
"root": {
"inputs": {
"fenix": "fenix",
"flake-utils": "flake-utils",
"nixpkgs": "nixpkgs"
}
},
"rust-analyzer-src": {
"flake": false,
"locked": {
"lastModified": 1717583671,
"narHash": "sha256-+lRAmz92CNUxorqWusgJbL9VE1eKCnQQojglRemzwkw=",
"owner": "rust-lang",
"repo": "rust-analyzer",
"rev": "48bbdd6a74f3176987d5c809894ac33957000d19",
"type": "github"
},
"original": {
"owner": "rust-lang",
"ref": "nightly",
"repo": "rust-analyzer",
"type": "github"
}
},
"systems": {
"locked": {
"lastModified": 1681028828,
"narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
"owner": "nix-systems",
"repo": "default",
"rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
"type": "github"
},
"original": {
"owner": "nix-systems",
"repo": "default",
"type": "github"
}
}
},
"root": "root",
"version": 7
}

60
flake.nix Normal file
View file

@ -0,0 +1,60 @@
{
inputs = {
nixpkgs.url = "github:nixos/nixpkgs/release-24.05";
flake-utils.url = "github:numtide/flake-utils";
fenix.url = "github:nix-community/fenix";
fenix.inputs.nixpkgs.follows = "nixpkgs";
};
outputs = { self, ... }@inputs:
inputs.flake-utils.lib.eachSystem
(with inputs.flake-utils.lib.system; [ x86_64-linux ])
(system:
let
pkgs = inputs.nixpkgs.legacyPackages.${system}.extend
inputs.fenix.overlays.default;
inherit (pkgs) lib;
in
{
devShell = pkgs.mkShell rec {
packages = with pkgs; [
(fenix.complete.withComponents [
"cargo"
"clippy"
"rust-src"
"rustc"
"rustfmt"
])
rust-analyzer-nightly
ruff
clang
llvmPackages.clang
pkg-config
leptonica
tesseract
openssl
sqlite
];
LD_LIBRARY_PATH = lib.makeLibraryPath packages;
# compilation of -sys packages requires manually setting LIBCLANG_PATH
LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib";
};
});
# {{{ Caching and whatnot
# TODO: persist trusted substituters file
nixConfig = {
extra-substituters = [
"https://nix-community.cachix.org"
];
extra-trusted-public-keys = [
"nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
];
};
# }}}
}

1
rustfmt.toml Normal file
View file

@ -0,0 +1 @@
hard_tabs=true

55
schema.sql Normal file
View file

@ -0,0 +1,55 @@
# {{{ users
create table IF NOT EXISTS users (
id INTEGER PRIMARY KEY,
discord_id TEXT UNIQUE NOT NULL,
nickname TEXT UNIQUE
);
# }}}
# {{{ songs
CREATE TABLE IF NOT EXISTS songs (
id INTEGER PRIMARY KEY,
title TEXT NOT NULL,
ocr_alias TEXT,
artist TEXT,
UNIQUE(title, artist)
);
# }}}
# {{{ charts
CREATE TABLE IF NOT EXISTS charts (
id INTEGER PRIMARY KEY,
song_id INTEGER NOT NULL,
difficulty TEXT NOT NULL CHECK (difficulty IN ('PST','PRS','FTR','ETR','BYD')),
level TEXT NOT NULL,
note_count INTEGER NOT NULL,
chart_constant INTEGER NOT NULL,
FOREIGN KEY (song_id) REFERENCES songs(id),
UNIQUE(song_id, difficulty)
);
# }}}
# {{{ plays
CREATE TABLE IF NOT EXISTS plays (
id INTEGER PRIMARY KEY,
chart_id INTEGER NOT NULL,
user_id INTEGER NOT NULL,
discord_attachment_id TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
creation_ptt INTEGER,
creation_zeta_ptt INTEGER,
score INTEGER NOT NULL,
zeta_score INTEGER,
max_recall INTEGER,
far_notes INTEGER,
FOREIGN KEY (chart_id) REFERENCES charts(id),
FOREIGN KEY (user_id) REFERENCES users(id)
);
# }}}
insert into users(discord_id, nickname) values (385759924917108740, 'prescientmoon');

73
scripts/main.py Executable file
View file

@ -0,0 +1,73 @@
#!/usr/bin/env nix-shell
#!nix-shell -p "pkgs.python3.withPackages (p: with p; [tabulate])"
#!nix-shell -i python3
import csv
import os
import sqlite3
import sys
data_dir = os.environ.get("SHIMMERING_DATA_DIR")
db_path = data_dir + "/db.sqlite"
# if not os.path.exists(db_path):
# run(f"cat ./schema.sql | sqlite3 {db_path}")
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()
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))
for title, charts in songs.items():
artist = None
if title.startswith("Quon"):
artist = title[6:-1]
title = "Quon"
row = 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(?,?,?,?,?)
""",
(
song_id,
difficulty,
level,
int(note_count.replace(",", "").replace(".", "")),
int(float(cc) * 100),
),
)
conn.commit()
print(f"Imported {chart_count} charts and {len(songs)} songs")
# }}}
command = sys.argv[1]
subcommand = sys.argv[2]
if command == "import" and subcommand == "charts":
import_charts_from_csv(sys.argv[3])

88
src/chart.rs Normal file
View file

@ -0,0 +1,88 @@
use sqlx::prelude::FromRow;
use crate::context::{Error, UserContext};
#[derive(Debug, Clone, Copy, sqlx::Type)]
pub enum Difficulty {
PST,
PRS,
FTR,
ETR,
BYD,
}
impl Difficulty {
#[inline]
pub fn to_index(self) -> usize {
self as usize
}
}
#[derive(Debug, Clone, FromRow)]
pub struct Song {
pub id: u32,
pub title: String,
pub ocr_alias: Option<String>,
pub artist: Option<String>,
}
#[derive(Debug, Clone, Copy, FromRow)]
pub struct Chart {
pub id: u32,
pub song_id: u32,
pub difficulty: Difficulty,
pub level: u32,
pub note_count: u32,
pub chart_constant: u32,
}
#[derive(Debug, Clone)]
pub struct CachedSong {
song: Song,
charts: [Option<Chart>; 5],
}
impl CachedSong {
pub fn new(song: Song, charts: [Option<Chart>; 5]) -> Self {
Self { song, charts }
}
}
#[derive(Debug, Clone, Default)]
pub struct SongCache {
songs: Vec<Option<CachedSong>>,
}
impl SongCache {
pub async fn new(ctx: &UserContext) -> Result<Self, Error> {
let mut result = Self::default();
let songs: Vec<Song> = sqlx::query_as("SELECT * FROM songs")
.fetch_all(&ctx.db)
.await?;
for song in songs {
let song_id = song.id as usize;
if song_id >= result.songs.len() {
result.songs.resize(song_id, None);
}
let charts: Vec<Chart> = sqlx::query_as("SELECT * FROM charts WHERE song_id=?")
.bind(song.id)
.fetch_all(&ctx.db)
.await?;
let mut chart_cache = [None; 5];
for chart in charts {
chart_cache[chart.difficulty.to_index()] = Some(chart);
}
result.songs[song_id] = Some(CachedSong::new(song, chart_cache));
}
Ok(result)
}
}

170
src/commands.rs Normal file
View file

@ -0,0 +1,170 @@
use crate::context::{Context, Error};
use crate::score::ImageCropper;
use crate::user::User;
use image::imageops::FilterType;
use poise::serenity_prelude::{
CreateAttachment, CreateEmbed, CreateEmbedAuthor, CreateMessage, Timestamp,
};
use poise::{serenity_prelude as serenity, CreateReply};
use prettytable::format::{FormatBuilder, LinePosition, LineSeparator};
use prettytable::{row, Table};
/// Show this help menu
#[poise::command(prefix_command, track_edits, slash_command)]
pub async fn help(
ctx: Context<'_>,
#[description = "Specific command to show help about"]
#[autocomplete = "poise::builtins::autocomplete_command"]
command: Option<String>,
) -> Result<(), Error> {
poise::builtins::help(
ctx,
command.as_deref(),
poise::builtins::HelpConfiguration {
extra_text_at_bottom: "For additional support, message @prescientmoon",
..Default::default()
},
)
.await?;
Ok(())
}
/// Score management
#[poise::command(
prefix_command,
slash_command,
subcommands("magic"),
subcommand_required
)]
pub async fn score(_ctx: Context<'_>) -> Result<(), Error> {
Ok(())
}
/// Identify scores from attached images.
#[poise::command(prefix_command, slash_command)]
pub async fn magic(
ctx: Context<'_>,
#[description = "Images containing scores"] files: Vec<serenity::Attachment>,
) -> Result<(), Error> {
println!("{:?}", User::from_context(&ctx).await);
if files.len() == 0 {
ctx.reply("No images found attached to message").await?;
} else {
let handle = ctx
.reply(format!("Processing: 0/{} images", files.len()))
.await?;
for (i, file) in files.iter().enumerate() {
if let Some(_) = file.dimensions() {
// Download image and guess it's format
let bytes = file.download().await?;
let format = image::guess_format(&bytes)?;
// Image pre-processing
let mut image = image::load_from_memory_with_format(&bytes, format)?
.resize(1024, 1024, FilterType::Nearest)
.grayscale()
.blur(1.);
image.invert();
// {{{ 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?;
// }}}
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?;
// Create cropper and run OCR
let mut cropper = ImageCropper::default();
let score_readout = match cropper.read_score(&image) {
// {{{ OCR error handling
Err(err) => {
let error_attachement =
CreateAttachment::bytes(cropper.bytes, &file.filename);
let msg = CreateMessage::default().embed(
CreateEmbed::default()
.title("Could not read score from picture")
.attachment(&file.filename)
.description(format!("{}", err))
.author(
CreateEmbedAuthor::new(&ctx.author().name)
.icon_url(ctx.author().face()),
)
.timestamp(Timestamp::now()),
);
ctx.channel_id()
.send_files(ctx.http(), [error_attachement], msg)
.await?;
continue;
}
// }}}
Ok(score) => score,
};
// Reply with attachement & readout
let attachement = CreateAttachment::bytes(cropper.bytes, &file.filename);
let reply = CreateReply::default()
.attachment(attachement)
.content(format!("Score: {:?}", score_readout))
.reply(true);
ctx.send(reply).await?;
// Edit progress reply
let progress_reply = CreateReply::default()
.content(format!("Processing: {}/{} images", i + 1, files.len()))
.reply(true);
handle.edit(ctx, progress_reply).await?;
} else {
ctx.reply("One of the attached files is not an image!")
.await?;
continue;
}
}
// Finish off progress reply
let progress_reply = CreateReply::default()
.content(format!("All images have been processed!"))
.reply(true);
handle.edit(ctx, progress_reply).await?;
}
Ok(())
}

23
src/context.rs Normal file
View file

@ -0,0 +1,23 @@
use sqlx::SqlitePool;
use crate::chart::SongCache;
// Types used by all command functions
pub type Error = Box<dyn std::error::Error + Send + Sync>;
pub type Context<'a> = poise::Context<'a, UserContext, Error>;
// Custom user data passed to all command functions
pub struct UserContext {
pub db: SqlitePool,
pub song_cache: SongCache,
}
impl UserContext {
#[inline]
pub fn new(db: SqlitePool) -> Self {
Self {
db,
song_cache: SongCache::default(),
}
}
}

86
src/main.rs Normal file
View file

@ -0,0 +1,86 @@
#![warn(clippy::str_to_string)]
#![feature(iter_map_windows)]
mod chart;
mod commands;
mod context;
mod score;
mod user;
use chart::SongCache;
use context::{Error, UserContext};
use poise::serenity_prelude as serenity;
use score::score_to_zeta_score;
use sqlx::sqlite::SqlitePoolOptions;
use std::{env::var, sync::Arc, time::Duration};
// {{{ Error handler
async fn on_error(error: poise::FrameworkError<'_, UserContext, Error>) {
match error {
poise::FrameworkError::Setup { error, .. } => panic!("Failed to start bot: {:?}", error),
poise::FrameworkError::Command { error, ctx, .. } => {
println!("Error in command `{}`: {:?}", ctx.command().name, error,);
}
error => {
if let Err(e) = poise::builtins::on_error(error).await {
println!("Error while handling error: {}", e)
}
}
}
}
// }}}
#[tokio::main]
async fn main() {
let data_dir = var("SHIMMERING_DATA_DIR")
.expect("Missing `SHIMMERING_DATA_DIR` env var, see README for more information.");
let pool = SqlitePoolOptions::new()
.connect(&format!("sqlite://{}/db.sqlite", data_dir))
.await
.unwrap();
println!("{:?}", score_to_zeta_score(9966677, 1303));
println!("{:?}", score_to_zeta_score(9970525, 1303));
// {{{ Poise options
let options = poise::FrameworkOptions {
commands: vec![commands::help(), commands::score()],
prefix_options: poise::PrefixFrameworkOptions {
prefix: Some("!".into()),
edit_tracker: Some(Arc::new(poise::EditTracker::for_timespan(
Duration::from_secs(3600),
))),
..Default::default()
},
// The global error handler for all error cases that may occur
on_error: |error| Box::pin(on_error(error)),
..Default::default()
};
// }}}
// {{{ Start poise
let framework = poise::Framework::builder()
.setup(move |ctx, _ready, framework| {
Box::pin(async move {
println!("Logged in as {}", _ready.user.name);
poise::builtins::register_globally(ctx, &framework.options().commands).await?;
let mut ctx = UserContext::new(pool);
ctx.song_cache = SongCache::new(&ctx).await?;
Ok(ctx)
})
})
.options(options)
.build();
let token = var("SHIMMERING_DISCORD_TOKEN")
.expect("Missing `SHIMMERING_DISCORD_TOKEN` env var, see README for more information.");
let intents =
serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::MESSAGE_CONTENT;
let client = serenity::ClientBuilder::new(token, intents)
.framework(framework)
.await;
client.unwrap().start().await.unwrap()
// }}}
}

385
src/score.rs Normal file
View file

@ -0,0 +1,385 @@
#![allow(dead_code)]
use std::{io::Cursor, sync::OnceLock, time::Instant};
use image::DynamicImage;
use num::Rational64;
use tesseract::{PageSegMode, Tesseract};
use crate::{
chart::{Chart, Difficulty},
context::{Error, UserContext},
user::User,
};
// {{{ ImageDimensions
#[derive(Debug, Clone, Copy)]
pub struct ImageDimensions {
width: u32,
height: u32,
}
impl ImageDimensions {
#[inline]
pub fn new(width: u32, height: u32) -> Self {
Self { width, height }
}
#[inline]
pub fn aspect_ratio(&self) -> f32 {
self.width as f32 / self.height as f32
}
#[inline]
pub fn from_image(image: &DynamicImage) -> Self {
Self::new(image.width(), image.height())
}
}
// }}}
// {{{ AbsoluteRect
#[derive(Debug, Clone, Copy)]
pub struct AbsoluteRect {
pub x: u32,
pub y: u32,
pub width: u32,
pub height: u32,
pub dimensions: ImageDimensions,
}
impl AbsoluteRect {
#[inline]
pub fn new(x: u32, y: u32, width: u32, height: u32, dimensions: ImageDimensions) -> Self {
Self {
x,
y,
width,
height,
dimensions,
}
}
#[inline]
pub fn to_relative(&self) -> RelativeRect {
RelativeRect::new(
self.x as f32 / self.dimensions.width as f32,
self.y as f32 / self.dimensions.height as f32,
self.width as f32 / self.dimensions.width as f32,
self.height as f32 / self.dimensions.height as f32,
self.dimensions,
)
}
}
// }}}
// {{{ RelativeRect
#[derive(Debug, Clone, Copy)]
pub struct RelativeRect {
pub x: f32,
pub y: f32,
pub width: f32,
pub height: f32,
pub dimensions: ImageDimensions,
}
fn lerp(i: f32, a: f32, b: f32) -> f32 {
a + (b - a) * i
}
impl RelativeRect {
#[inline]
pub fn new(x: f32, y: f32, width: f32, height: f32, dimensions: ImageDimensions) -> Self {
Self {
x,
y,
width,
height,
dimensions,
}
}
#[inline]
pub fn to_absolute(&self) -> AbsoluteRect {
AbsoluteRect::new(
(self.x * self.dimensions.width as f32) as u32,
(self.y * self.dimensions.height as f32) as u32,
(self.width * self.dimensions.width as f32) as u32,
(self.height * self.dimensions.height as f32) as u32,
self.dimensions,
)
}
pub fn from_aspect_ratio(
dimensions: ImageDimensions,
datapoints: &[RelativeRect],
) -> Option<Self> {
let aspect_ratio = dimensions.aspect_ratio();
for i in 0..(datapoints.len() - 1) {
let low = datapoints[i];
let high = datapoints[i + 1];
let low_ratio = low.dimensions.aspect_ratio();
let high_ratio = high.dimensions.aspect_ratio();
if (i == 0 || low_ratio <= aspect_ratio)
&& (aspect_ratio <= high_ratio || i == datapoints.len() - 2)
{
let p = (aspect_ratio - low_ratio) / (high_ratio - low_ratio);
return Some(Self::new(
lerp(p, low.x, high.x),
lerp(p, low.y, high.y) - 0.005,
lerp(p, low.width, high.width),
lerp(p, low.height, high.height) + 2. * 0.005,
dimensions,
));
}
}
None
}
}
// }}}
// {{{ Data points
fn score_rects() -> &'static [RelativeRect] {
static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
CELL.get_or_init(|| {
let mut rects: Vec<RelativeRect> = vec![
AbsoluteRect::new(642, 287, 284, 51, ImageDimensions::new(1560, 720)).to_relative(),
AbsoluteRect::new(651, 285, 305, 55, ImageDimensions::new(1600, 720)).to_relative(),
AbsoluteRect::new(748, 485, 503, 82, ImageDimensions::new(2000, 1200)).to_relative(),
AbsoluteRect::new(841, 683, 500, 92, ImageDimensions::new(2160, 1620)).to_relative(),
AbsoluteRect::new(851, 707, 532, 91, ImageDimensions::new(2224, 1668)).to_relative(),
AbsoluteRect::new(1037, 462, 476, 89, ImageDimensions::new(2532, 1170)).to_relative(),
AbsoluteRect::new(973, 653, 620, 105, ImageDimensions::new(2560, 1600)).to_relative(),
AbsoluteRect::new(1069, 868, 636, 112, ImageDimensions::new(2732, 2048)).to_relative(),
AbsoluteRect::new(1125, 510, 534, 93, ImageDimensions::new(2778, 1284)).to_relative(),
];
rects.sort_by_key(|r| (r.dimensions.aspect_ratio() * 1000.0).floor() as u32);
// Filter datapoints that are close together
let mut i = 0;
while i < rects.len() - 1 {
let low = rects[i];
let high = rects[i + 1];
if (low.dimensions.aspect_ratio() - high.dimensions.aspect_ratio()).abs() < 0.001 {
// TODO: we could interpolate here but oh well
rects.remove(i + 1);
}
i += 1;
}
rects
})
}
fn difficulty_rects() -> &'static [RelativeRect] {
static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
CELL.get_or_init(|| {
let mut rects: Vec<RelativeRect> = vec![
AbsoluteRect::new(642, 287, 284, 51, ImageDimensions::new(1560, 720)).to_relative(),
AbsoluteRect::new(651, 285, 305, 55, ImageDimensions::new(1600, 720)).to_relative(),
AbsoluteRect::new(748, 485, 503, 82, ImageDimensions::new(2000, 1200)).to_relative(),
AbsoluteRect::new(841, 683, 500, 92, ImageDimensions::new(2160, 1620)).to_relative(),
AbsoluteRect::new(851, 707, 532, 91, ImageDimensions::new(2224, 1668)).to_relative(),
AbsoluteRect::new(1037, 462, 476, 89, ImageDimensions::new(2532, 1170)).to_relative(),
AbsoluteRect::new(973, 653, 620, 105, ImageDimensions::new(2560, 1600)).to_relative(),
AbsoluteRect::new(1069, 868, 636, 112, ImageDimensions::new(2732, 2048)).to_relative(),
AbsoluteRect::new(1125, 510, 534, 93, ImageDimensions::new(2778, 1284)).to_relative(),
];
rects.sort_by_key(|r| (r.dimensions.aspect_ratio() * 1000.0).floor() as u32);
rects
})
}
// }}}
// {{{ Plays
/// Returns the zeta score and the number of shinies
pub fn score_to_zeta_score(score: u32, note_count: u32) -> (u32, 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 score = Rational64::from_integer(score as i64);
let score_units = (score / increment).floor();
let non_shiny_score = (score_units * increment).floor();
let shinies = score - non_shiny_score;
let zeta_score_units = Rational64::from_integer(2) * score_units + shinies;
let zeta_score = (zeta_increment * zeta_score_units).floor().to_integer() as u32;
(zeta_score, shinies.to_integer() as u32)
}
// {{{ Create play
#[derive(Debug, Clone)]
pub struct CreatePlay {
chart_id: u32,
user_id: u32,
discord_attachment_id: Option<String>,
// Actual score data
score: u32,
zeta_score: Option<u32>,
// Optional score details
max_recall: Option<u32>,
far_notes: Option<u32>,
// Creation data
creation_ptt: Option<u32>,
creation_zeta_ptt: Option<u32>,
}
impl CreatePlay {
#[inline]
pub fn new(score: u32, chart: Chart, user: User) -> Self {
Self {
chart_id: chart.id,
user_id: user.id,
discord_attachment_id: None,
score,
zeta_score: Some(score_to_zeta_score(score, chart.note_count).0),
max_recall: None,
far_notes: None,
// TODO: populate these
creation_ptt: None,
creation_zeta_ptt: None,
}
}
pub async fn save(self, ctx: &UserContext) -> Result<Play, Error> {
let play = sqlx::query_as!(
Play,
"
INSERT INTO plays(
user_id,chart_id,discord_attachment_id,
score,zeta_score,max_recall,far_notes
)
VALUES(?,?,?,?,?,?,?)
RETURNING *
",
self.user_id,
self.chart_id,
self.discord_attachment_id,
self.score,
self.zeta_score,
self.max_recall,
self.far_notes
)
.fetch_one(&ctx.db)
.await?;
Ok(play)
}
}
// }}}
// {{{ Play
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct Play {
id: i64,
chart_id: i64,
user_id: i64,
discord_attachment_id: Option<String>,
// Actual score data
score: i64,
zeta_score: Option<i64>,
// Optional score details
max_recall: Option<i64>,
far_notes: Option<i64>,
// Creation data
created_at: chrono::NaiveDateTime,
creation_ptt: Option<i64>,
creation_zeta_ptt: Option<i64>,
}
// }}}
// {{{ Tests
#[cfg(test)]
mod score_tests {
use super::*;
#[test]
fn zeta_score_consistent_with_pms() {
// note counts
for note_count in 200..=2000 {
for shiny_count in 0..=note_count {
let score = 10000000 + shiny_count;
let zeta_score_units = 4 * (note_count - shiny_count) + 5 * shiny_count;
let expected_zeta_score = Rational64::from_integer(zeta_score_units as i64)
* Rational64::new_raw(2000000, note_count as i64).reduced();
let (zeta_score, computed_shiny_count) = score_to_zeta_score(score, note_count);
assert_eq!(zeta_score, expected_zeta_score.to_integer() as u32);
assert_eq!(computed_shiny_count, shiny_count);
}
}
}
}
// }}}
// }}}
// {{{ Ocr types
#[derive(Debug, Clone, Copy)]
pub struct ScoreReadout {
pub score: u32,
pub difficulty: Difficulty,
}
impl ScoreReadout {
pub fn new(score: u32, difficulty: Difficulty) -> Self {
Self { score, difficulty }
}
}
// }}}
// {{{ Run OCR
/// Caches a byte vector in order to prevent reallocation
#[derive(Debug, Clone, Default)]
pub struct ImageCropper {
/// cached byte array
pub bytes: Vec<u8>,
}
impl ImageCropper {
fn crop_image_to_bytes(
&mut self,
image: &DynamicImage,
rect: AbsoluteRect,
) -> Result<(), Error> {
self.bytes.clear();
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)?;
Ok(())
}
pub fn read_score(&mut self, image: &DynamicImage) -> Result<ScoreReadout, Error> {
let rect =
RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), score_rects())
.ok_or_else(|| "Could not find score area in picture")?
.to_absolute();
self.crop_image_to_bytes(&image, rect)?;
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 = t.recognize()?;
if t.mean_text_conf() < 10 {
Err("Score text is not readable.")?;
}
let text: String = t
.get_text()?
.trim()
.chars()
.filter(|char| *char != ' ' && *char != '\'')
.collect();
let int = u32::from_str_radix(&text, 10)?;
Ok(ScoreReadout::new(int, Difficulty::FTR))
}
}
// }}}

20
src/user.rs Normal file
View file

@ -0,0 +1,20 @@
use crate::context::{Context, Error};
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct User {
pub id: u32,
pub discord_id: String,
pub nickname: Option<String>,
}
impl User {
pub async fn from_context(ctx: &Context<'_>) -> Result<Self, Error> {
let id = ctx.author().id.get().to_string();
let user = sqlx::query_as("SELECT * FROM users WHERE discord_id = ?")
.bind(id)
.fetch_one(&ctx.data().db)
.await?;
Ok(user)
}
}