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
|
.env
|
||||||
config.toml
|
config.toml
|
||||||
/target
|
/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 = [
|
dependencies = [
|
||||||
"dotenvy",
|
"dotenvy",
|
||||||
"image",
|
"image",
|
||||||
|
"once_cell",
|
||||||
"serde",
|
"serde",
|
||||||
"serenity",
|
"serenity",
|
||||||
"swordfish-common",
|
"swordfish-common",
|
||||||
|
@ -5,13 +5,11 @@ pub mod constants;
|
|||||||
pub mod tesseract;
|
pub mod tesseract;
|
||||||
|
|
||||||
pub fn setup_logger(level: &str) -> Result<(), fern::InitError> {
|
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()
|
let formatter = fmt::format()
|
||||||
.with_level(true)
|
.with_level(true)
|
||||||
.with_target(true)
|
.with_target(true)
|
||||||
.with_thread_ids(false)
|
.with_thread_ids(false)
|
||||||
.with_thread_names(false); // include the name of the current thread.pretty();
|
.with_thread_names(false);
|
||||||
let filter = EnvFilter::builder()
|
let filter = EnvFilter::builder()
|
||||||
.from_env()
|
.from_env()
|
||||||
.unwrap()
|
.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]
|
[dependencies]
|
||||||
dotenvy = "0.15.7"
|
dotenvy = "0.15.7"
|
||||||
image = "0.24.7"
|
image = "0.24.7"
|
||||||
|
once_cell = "1.19.0"
|
||||||
serde = "1.0.193"
|
serde = "1.0.193"
|
||||||
serenity = "0.12.0"
|
serenity = { version = "0.12.0", features = ["builder"] }
|
||||||
tokio = { version = "1.35.1", features = ["full"] }
|
tokio = { version = "1.35.1", features = ["full"] }
|
||||||
toml = "0.8.8"
|
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::model::channel::Message;
|
||||||
use serenity::prelude::*;
|
use serenity::prelude::*;
|
||||||
use std::io::Cursor;
|
use std::io::Cursor;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::thread;
|
||||||
use swordfish_common::tesseract::LepTess;
|
use swordfish_common::tesseract::LepTess;
|
||||||
use swordfish_common::{debug, error, info, trace, warn};
|
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 {
|
if message.attachments.len() < 1 {
|
||||||
return Err("No attachments found".to_string());
|
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
|
// Get the image attachment
|
||||||
let attachment = &message.attachments[0];
|
let attachment = &message.attachments[0];
|
||||||
let image_bytes = match attachment.download().await {
|
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)),
|
Err(why) => return Err(format!("Failed to read image: {:?}", why)),
|
||||||
};
|
};
|
||||||
|
trace!("Grayscaling image...");
|
||||||
img = img.grayscale();
|
img = img.grayscale();
|
||||||
img.save("debug.png").unwrap();
|
img.save("debug/1-grayscale.png").unwrap();
|
||||||
match lep_tess.set_image_from_mem(&img.as_bytes()) {
|
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(_) => (),
|
Ok(_) => (),
|
||||||
Err(why) => return Err(format!("Failed to set image: {:?}", why)),
|
Err(why) => return Err(format!("Failed to set image: {:?}", why)),
|
||||||
};
|
};
|
||||||
|
@ -1,4 +1,5 @@
|
|||||||
use dotenvy::dotenv;
|
use dotenvy::dotenv;
|
||||||
|
use once_cell::sync::Lazy;
|
||||||
use serenity::async_trait;
|
use serenity::async_trait;
|
||||||
use serenity::framework::standard::macros::{command, group};
|
use serenity::framework::standard::macros::{command, group};
|
||||||
use serenity::framework::standard::{CommandResult, Configuration, StandardFramework};
|
use serenity::framework::standard::{CommandResult, Configuration, StandardFramework};
|
||||||
@ -9,22 +10,30 @@ use serenity::model::{
|
|||||||
use serenity::prelude::*;
|
use serenity::prelude::*;
|
||||||
use std::env;
|
use std::env;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
use swordfish_common::*;
|
use swordfish_common::*;
|
||||||
|
|
||||||
mod config;
|
mod config;
|
||||||
|
mod helper;
|
||||||
mod katana;
|
mod katana;
|
||||||
|
mod template;
|
||||||
|
|
||||||
const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish";
|
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]
|
#[group]
|
||||||
#[commands(ping, drop_analyze)]
|
#[commands(ping, kdropanalyze)]
|
||||||
struct General;
|
struct General;
|
||||||
struct Handler;
|
struct Handler;
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl EventHandler for Handler {
|
impl EventHandler for Handler {
|
||||||
async fn message(&self, ctx: Context, msg: Message) {
|
async fn message(&self, ctx: Context, msg: Message) {
|
||||||
if msg.author.id == ctx.cache.current_user().id {
|
if msg.author.id == ctx.cache.current_user().id {
|
||||||
trace!("Ignoring message from self");
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
trace!("Message: {}, sender: {}", msg.content, msg.author.id);
|
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!")
|
.contains("I'm dropping 3 cards since this server is currently active!")
|
||||||
{
|
{
|
||||||
trace!("Card drop detected, executing drop analyzer...");
|
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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@ -95,32 +114,31 @@ async fn ping(ctx: &Context, msg: &Message) -> CommandResult {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[command]
|
#[command]
|
||||||
async fn drop_analyze(ctx: &Context, msg: &Message) -> CommandResult {
|
async fn kdropanalyze(ctx: &Context, msg: &Message) -> CommandResult {
|
||||||
let target_channel_id = match msg.content.split(" ").nth(1) {
|
let mut args = msg.content.split(" ");
|
||||||
|
let target_channel_id = match args.nth(1) {
|
||||||
Some(content) => match content.parse::<u64>() {
|
Some(content) => match content.parse::<u64>() {
|
||||||
Ok(id) => id,
|
Ok(id) => id,
|
||||||
Err(why) => {
|
Err(why) => {
|
||||||
msg.reply(ctx, format!("Failed to parse message ID: {:?}", why))
|
helper::error_message(ctx, msg, format!("Failed to parse channel ID: `{:?}`", why)).await;
|
||||||
.await?;
|
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
None => {
|
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(());
|
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>() {
|
Some(content) => match content.parse::<u64>() {
|
||||||
Ok(id) => id,
|
Ok(id) => id,
|
||||||
Err(why) => {
|
Err(why) => {
|
||||||
msg.reply(ctx, format!("Failed to parse message ID: {:?}", why))
|
helper::error_message(ctx, msg, format!("Failed to parse message ID: `{:?}`", why)).await;
|
||||||
.await?;
|
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
None => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -134,19 +152,19 @@ async fn drop_analyze(ctx: &Context, msg: &Message) -> CommandResult {
|
|||||||
{
|
{
|
||||||
Ok(msg) => msg,
|
Ok(msg) => msg,
|
||||||
Err(why) => {
|
Err(why) => {
|
||||||
msg.reply(ctx, format!("Failed to get message: {:?}", why))
|
helper::error_message(ctx, msg, format!("Failed to get message: `{:?}`", why)).await;
|
||||||
.await?;
|
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
match katana::analyze_drop_message(&target_msg).await {
|
unsafe {
|
||||||
Ok(_) => {
|
match katana::analyze_drop_message(&LEPTESS_ARC, &target_msg).await {
|
||||||
msg.reply(ctx, "Drop analysis complete").await?;
|
Ok(_) => {
|
||||||
}
|
msg.reply(ctx, "Drop analysis complete").await?;
|
||||||
Err(why) => {
|
}
|
||||||
msg.reply(ctx, format!("Failed to analyze drop: {:?}", why))
|
Err(why) => {
|
||||||
.await?;
|
helper::error_message(ctx, msg, format!("Failed to analyze drop: `{:?}`", why)).await;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
}
|
||||||
Ok(())
|
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