From 34c52cd9ab60b3673d77771ae8f1f5c384a5d610 Mon Sep 17 00:00:00 2001 From: tretrauit Date: Sun, 31 Dec 2023 22:55:32 +0700 Subject: [PATCH] feat: use Arc & Mutex for LepTess --- .gitignore | 1 + .vscode/settings.json | 6 ++ Cargo.lock | 1 + swordfish-common/src/lib.rs | 4 +- swordfish-common/src/tesseract.rs | 14 +++- .../src/utils/{karuta.rs => katana.rs} | 0 swordfish/Cargo.toml | 3 +- swordfish/src/helper.rs | 20 ++++++ swordfish/src/katana.rs | 67 +++++++++++++++++-- swordfish/src/main.rs | 64 +++++++++++------- swordfish/src/template/message.rs | 43 ++++++++++++ swordfish/src/template/mod.rs | 1 + 12 files changed, 191 insertions(+), 33 deletions(-) create mode 100644 .vscode/settings.json rename swordfish-common/src/utils/{karuta.rs => katana.rs} (100%) create mode 100644 swordfish/src/helper.rs create mode 100644 swordfish/src/template/message.rs create mode 100644 swordfish/src/template/mod.rs diff --git a/.gitignore b/.gitignore index bbed3f6..776e7e1 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ .env config.toml /target +/debug \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..6633290 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,6 @@ +{ + "rust-analyzer.linkedProjects": [ + ".\\swordfish\\Cargo.toml" + ], + "rust-analyzer.showUnlinkedFileNotification": false +} \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index f875f68..47510b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1682,6 +1682,7 @@ version = "0.1.0" dependencies = [ "dotenvy", "image", + "once_cell", "serde", "serenity", "swordfish-common", diff --git a/swordfish-common/src/lib.rs b/swordfish-common/src/lib.rs index 979a793..fd1bc2c 100644 --- a/swordfish-common/src/lib.rs +++ b/swordfish-common/src/lib.rs @@ -5,13 +5,11 @@ pub mod constants; pub mod tesseract; pub fn setup_logger(level: &str) -> Result<(), fern::InitError> { - // I don't really know how to do it because the unset variable trick doesn't work - // since the types can be let formatter = fmt::format() .with_level(true) .with_target(true) .with_thread_ids(false) - .with_thread_names(false); // include the name of the current thread.pretty(); + .with_thread_names(false); let filter = EnvFilter::builder() .from_env() .unwrap() diff --git a/swordfish-common/src/tesseract.rs b/swordfish-common/src/tesseract.rs index a416b0b..eccf7aa 100644 --- a/swordfish-common/src/tesseract.rs +++ b/swordfish-common/src/tesseract.rs @@ -1 +1,13 @@ -pub use leptess::LepTess; +pub use leptess::{LepTess, Variable}; + +pub fn init_tesseract() -> Result { + let mut lep_tess = match LepTess::new(None, "eng") { + Ok(lep_tess) => lep_tess, + Err(why) => return Err(format!("Failed to initialize Tesseract: {:?}", why)), + }; + match lep_tess.set_variable(Variable::TesseditCharWhitelist, "0123456789") { + Ok(_) => (), + Err(why) => return Err(format!("Failed to set whitelist: {:?}", why)), + }; + Ok(lep_tess) +} diff --git a/swordfish-common/src/utils/karuta.rs b/swordfish-common/src/utils/katana.rs similarity index 100% rename from swordfish-common/src/utils/karuta.rs rename to swordfish-common/src/utils/katana.rs diff --git a/swordfish/Cargo.toml b/swordfish/Cargo.toml index 60f2018..fb70ff2 100644 --- a/swordfish/Cargo.toml +++ b/swordfish/Cargo.toml @@ -8,8 +8,9 @@ edition = "2021" [dependencies] dotenvy = "0.15.7" image = "0.24.7" +once_cell = "1.19.0" serde = "1.0.193" -serenity = "0.12.0" +serenity = { version = "0.12.0", features = ["builder"] } tokio = { version = "1.35.1", features = ["full"] } toml = "0.8.8" diff --git a/swordfish/src/helper.rs b/swordfish/src/helper.rs new file mode 100644 index 0000000..fbc9fed --- /dev/null +++ b/swordfish/src/helper.rs @@ -0,0 +1,20 @@ +use serenity::client::Context; +use serenity::model::channel::Message; +use serenity::builder::CreateMessage; +use crate::template::message; + +pub async fn error_message(ctx: &Context, msg: &Message, content: String) { + msg.channel_id + .send_message( + ctx, + CreateMessage::new().add_embed( + message::error_embed( + ctx, + None, + Some(content), + ) + .await, + ), + ) + .await?; +} \ No newline at end of file diff --git a/swordfish/src/katana.rs b/swordfish/src/katana.rs index 1bcfff1..06267f1 100644 --- a/swordfish/src/katana.rs +++ b/swordfish/src/katana.rs @@ -4,15 +4,25 @@ use serenity::framework::standard::{CommandResult, Configuration, StandardFramew use serenity::model::channel::Message; use serenity::prelude::*; use std::io::Cursor; +use std::sync::{Arc, Mutex}; +use std::thread; use swordfish_common::tesseract::LepTess; use swordfish_common::{debug, error, info, trace, warn}; -pub async fn analyze_drop_message(message: &Message) -> Result<(), String> { +pub fn analyze_card(leptess: &LepTess, card: image::DynamicImage) { + trace!("Analyzing card..."); + // Read the name and the series + let name_img = card.crop_imm(0, 0, 257, 29); + // Read the print number +} + +pub async fn analyze_drop_message( + leptess_arc: &Arc>, + message: &Message, +) -> Result<(), String> { if message.attachments.len() < 1 { return Err("No attachments found".to_string()); }; - trace!("Initializing Tesseract OCR engine..."); - let mut lep_tess = LepTess::new(None, "eng").unwrap(); // Get the image attachment let attachment = &message.attachments[0]; let image_bytes = match attachment.download().await { @@ -27,9 +37,56 @@ pub async fn analyze_drop_message(message: &Message) -> Result<(), String> { }, Err(why) => return Err(format!("Failed to read image: {:?}", why)), }; + trace!("Grayscaling image..."); img = img.grayscale(); - img.save("debug.png").unwrap(); - match lep_tess.set_image_from_mem(&img.as_bytes()) { + img.save("debug/1-grayscale.png").unwrap(); + trace!("Increasing contrast of the image..."); + img = img.adjust_contrast(1.0); + img.save("debug/2-contrast.png").unwrap(); + // Cropping cards + let distance = 257 - 29 + 305 - 259; + let cards_count = img.width() / distance; + trace!("Cropping {} cards...", cards_count); + let mut jobs: Vec<_> = Vec::new(); + for i_real in 0..cards_count { + let i = i_real.clone(); + let leptess_mutex = leptess_arc.clone(); + let img = img.clone(); + let job = move || { + Ok({ + let x = 29 + distance * i; + let y = 34; + let width = 257 + distance * i - x; + let height = 387 - y; + trace!("Cropping card {} ({}, {}, {}, {})", i, x, y, width, height); + let card_img = img.crop_imm(x, y, width, height); + match card_img.save(format!("debug/3-cropped-{}.png", i)) { + Ok(_) => { + trace!("Saved cropped card {}", i); + let leptess = leptess_mutex.lock().unwrap(); + analyze_card(&leptess, card_img); + } + Err(why) => return Err(format!("Failed to save image: {:?}", why)), + }; + }) + }; + jobs.push(job); + } + let mut tasks: Vec>> = Vec::new(); + for job in jobs { + let task = thread::spawn(job); + tasks.push(task); + } + for task in tasks { + let result = task.join(); + match result { + Ok(_) => (), + Err(why) => return Err(format!("Failed to crop card: {:?}", why)), + }; + } + let leptess_mutex = leptess_arc.clone(); + let mut leptess = leptess_mutex.lock().unwrap(); + match leptess.set_image_from_mem(&img.as_bytes()) { Ok(_) => (), Err(why) => return Err(format!("Failed to set image: {:?}", why)), }; diff --git a/swordfish/src/main.rs b/swordfish/src/main.rs index 4470c8a..453b424 100644 --- a/swordfish/src/main.rs +++ b/swordfish/src/main.rs @@ -1,4 +1,5 @@ use dotenvy::dotenv; +use once_cell::sync::Lazy; use serenity::async_trait; use serenity::framework::standard::macros::{command, group}; use serenity::framework::standard::{CommandResult, Configuration, StandardFramework}; @@ -9,22 +10,30 @@ use serenity::model::{ use serenity::prelude::*; use std::env; use std::path::Path; +use std::sync::{Arc, Mutex}; use swordfish_common::*; mod config; +mod helper; mod katana; +mod template; const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish"; +static mut LEPTESS_ARC: Lazy>> = Lazy::new(|| { + trace!("Initializing Tesseract..."); + Arc::new(Mutex::new( + tesseract::init_tesseract().expect("Failed to initialize Tesseract"), + )) +}); #[group] -#[commands(ping, drop_analyze)] +#[commands(ping, kdropanalyze)] struct General; struct Handler; #[async_trait] impl EventHandler for Handler { async fn message(&self, ctx: Context, msg: Message) { if msg.author.id == ctx.cache.current_user().id { - trace!("Ignoring message from self"); return; } trace!("Message: {}, sender: {}", msg.content, msg.author.id); @@ -47,7 +56,17 @@ async fn parse_katana(ctx: &Context, msg: &Message) -> Result<(), String> { .contains("I'm dropping 3 cards since this server is currently active!") { trace!("Card drop detected, executing drop analyzer..."); - katana::analyze_drop_message(msg).await?; + unsafe { + match katana::analyze_drop_message(&LEPTESS_ARC, msg) { + Ok(_) => { + // msg.reply(ctx, "Drop analysis complete").await?; + } + Err(why) => { + trace!("Failed to analyze drop: `{:?}`", why); + // helper::error_message(ctx, msg, format!("Failed to analyze drop: `{:?}`", why)).await; + } + }; + } } Ok(()) } @@ -95,32 +114,31 @@ async fn ping(ctx: &Context, msg: &Message) -> CommandResult { } #[command] -async fn drop_analyze(ctx: &Context, msg: &Message) -> CommandResult { - let target_channel_id = match msg.content.split(" ").nth(1) { +async fn kdropanalyze(ctx: &Context, msg: &Message) -> CommandResult { + let mut args = msg.content.split(" "); + let target_channel_id = match args.nth(1) { Some(content) => match content.parse::() { Ok(id) => id, Err(why) => { - msg.reply(ctx, format!("Failed to parse message ID: {:?}", why)) - .await?; + helper::error_message(ctx, msg, format!("Failed to parse channel ID: `{:?}`", why)).await; return Ok(()); } }, None => { - msg.reply(ctx, "No message ID provided").await?; + helper::error_message(ctx, msg, "Usage: `kdropanalyze `".to_string()).await; return Ok(()); } }; - let target_msg_id = match msg.content.split(" ").nth(2) { + let target_msg_id = match args.nth(0) { Some(content) => match content.parse::() { Ok(id) => id, Err(why) => { - msg.reply(ctx, format!("Failed to parse message ID: {:?}", why)) - .await?; + helper::error_message(ctx, msg, format!("Failed to parse message ID: `{:?}`", why)).await; return Ok(()); } }, None => { - msg.reply(ctx, "No message ID provided").await?; + helper::error_message(ctx, msg, "Usage: `kdropanalyze `".to_string()).await; return Ok(()); } }; @@ -134,19 +152,19 @@ async fn drop_analyze(ctx: &Context, msg: &Message) -> CommandResult { { Ok(msg) => msg, Err(why) => { - msg.reply(ctx, format!("Failed to get message: {:?}", why)) - .await?; + helper::error_message(ctx, msg, format!("Failed to get message: `{:?}`", why)).await; return Ok(()); } }; - match katana::analyze_drop_message(&target_msg).await { - Ok(_) => { - msg.reply(ctx, "Drop analysis complete").await?; - } - Err(why) => { - msg.reply(ctx, format!("Failed to analyze drop: {:?}", why)) - .await?; - } - }; + unsafe { + match katana::analyze_drop_message(&LEPTESS_ARC, &target_msg).await { + Ok(_) => { + msg.reply(ctx, "Drop analysis complete").await?; + } + Err(why) => { + helper::error_message(ctx, msg, format!("Failed to analyze drop: `{:?}`", why)).await; + } + }; + } Ok(()) } diff --git a/swordfish/src/template/message.rs b/swordfish/src/template/message.rs new file mode 100644 index 0000000..c7e9d88 --- /dev/null +++ b/swordfish/src/template/message.rs @@ -0,0 +1,43 @@ +use serenity::builder::{CreateEmbed, CreateEmbedFooter}; +use serenity::client::Context; +use serenity::model::Color; + +pub async fn crate_embed( + client: &Context, + title: Option, + description: Option, + color: Color, +) -> CreateEmbed { + let user = client.http.get_current_user().await.unwrap(); + let embed = CreateEmbed::new() + .title(title.unwrap_or("Swordfish".to_string())) + .description(description.unwrap_or("".to_string())) + .color(color) + .footer( + CreateEmbedFooter::new(user.name.clone()) + .icon_url(user.avatar_url().unwrap_or("".to_string())), + ); + return embed; +} + +pub async fn error_embed( + client: &Context, + mut title: Option, + description: Option, +) -> CreateEmbed { + if title.is_none() { + title = Some("Error".to_string()); + } + return crate_embed(client, title, description, Color::RED).await; +} + +pub async fn info_embed( + client: &Context, + mut title: Option, + description: Option, +) -> CreateEmbed { + if title.is_none() { + title = Some("Info".to_string()); + } + return crate_embed(client, title, description, Color::DARK_GREEN).await; +} diff --git a/swordfish/src/template/mod.rs b/swordfish/src/template/mod.rs new file mode 100644 index 0000000..90d43c9 --- /dev/null +++ b/swordfish/src/template/mod.rs @@ -0,0 +1 @@ +pub mod message;