feat: use Arc & Mutex for LepTess
This commit is contained in:
parent
19dc2963f5
commit
34c52cd9ab
1
.gitignore
vendored
1
.gitignore
vendored
@ -2,3 +2,4 @@
|
||||
.env
|
||||
config.toml
|
||||
/target
|
||||
/debug
|
6
.vscode/settings.json
vendored
Normal file
6
.vscode/settings.json
vendored
Normal file
@ -0,0 +1,6 @@
|
||||
{
|
||||
"rust-analyzer.linkedProjects": [
|
||||
".\\swordfish\\Cargo.toml"
|
||||
],
|
||||
"rust-analyzer.showUnlinkedFileNotification": false
|
||||
}
|
1
Cargo.lock
generated
1
Cargo.lock
generated
@ -1682,6 +1682,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"dotenvy",
|
||||
"image",
|
||||
"once_cell",
|
||||
"serde",
|
||||
"serenity",
|
||||
"swordfish-common",
|
||||
|
@ -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()
|
||||
|
@ -1 +1,13 @@
|
||||
pub use leptess::LepTess;
|
||||
pub use leptess::{LepTess, Variable};
|
||||
|
||||
pub fn init_tesseract() -> Result<LepTess, String> {
|
||||
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)
|
||||
}
|
||||
|
@ -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"
|
||||
|
||||
|
20
swordfish/src/helper.rs
Normal file
20
swordfish/src/helper.rs
Normal file
@ -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?;
|
||||
}
|
@ -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<Mutex<LepTess>>,
|
||||
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<thread::JoinHandle<Result<(), String>>> = 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)),
|
||||
};
|
||||
|
@ -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<Arc<Mutex<tesseract::LepTess>>> = 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::<u64>() {
|
||||
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 <channel ID> <message ID>`".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::<u64>() {
|
||||
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 <channel ID> <message ID>`".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 {
|
||||
unsafe {
|
||||
match katana::analyze_drop_message(&LEPTESS_ARC, &target_msg).await {
|
||||
Ok(_) => {
|
||||
msg.reply(ctx, "Drop analysis complete").await?;
|
||||
}
|
||||
Err(why) => {
|
||||
msg.reply(ctx, format!("Failed to analyze drop: {:?}", why))
|
||||
.await?;
|
||||
helper::error_message(ctx, msg, format!("Failed to analyze drop: `{:?}`", why)).await;
|
||||
}
|
||||
};
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
43
swordfish/src/template/message.rs
Normal file
43
swordfish/src/template/message.rs
Normal file
@ -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<String>,
|
||||
description: Option<String>,
|
||||
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<String>,
|
||||
description: Option<String>,
|
||||
) -> 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<String>,
|
||||
description: Option<String>,
|
||||
) -> CreateEmbed {
|
||||
if title.is_none() {
|
||||
title = Some("Info".to_string());
|
||||
}
|
||||
return crate_embed(client, title, description, Color::DARK_GREEN).await;
|
||||
}
|
1
swordfish/src/template/mod.rs
Normal file
1
swordfish/src/template/mod.rs
Normal file
@ -0,0 +1 @@
|
||||
pub mod message;
|
Loading…
Reference in New Issue
Block a user