tretrauit/swordfish
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(katana): improve the image analyzing process

Browse Source 
Zamn, decreasing contrast and also copying Nori code.
This commit is contained in:
tretrauit
2024-01-10 01:24:22 +07:00
parent 1b943f5698
commit 6f35d05a3e
6 changed files with 79 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@ -2220,7 +2220,6 @@ dependencies = [
"dotenvy", "dotenvy",
"image", "image",
"leptess", "leptess",
"regex",
"rusty-tesseract", "rusty-tesseract",
"serde", "serde",
"serenity 0.12.0", "serenity 0.12.0",

4
swordfish-common/src/database/katana.rs
View File

@ -1,12 +1,12 @@
use crate::database; use crate::database;
use crate::structs::Character; use crate::structs::Character;
use mongodb::Collection; use mongodb::Collection;
use std::sync::OnceLock;
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::OnceCell;
use tokio::task; use tokio::task;
use tracing::trace; use tracing::trace;
pub static KATANA: OnceLock<Collection<Character>> = OnceLock::new(); pub static KATANA: OnceCell<Collection<Character>> = OnceCell::const_new();
/// ///
/// Initialize the "katana" collection in MongoDB /// Initialize the "katana" collection in MongoDB

6
swordfish-common/src/database/mod.rs
View File

@ -4,11 +4,11 @@ use mongodb::bson::doc;
use mongodb::options::ClientOptions; use mongodb::options::ClientOptions;
use mongodb::{Client, Database}; use mongodb::{Client, Database};
use std::env; use std::env;
use std::sync::OnceLock; use tokio::sync::OnceCell;
use tracing::info; use tracing::info;
static MONGO_CLIENT: OnceLock<Client> = OnceLock::new(); static MONGO_CLIENT: OnceCell<Client> = OnceCell::const_new();
static MONGO_DATABASE: OnceLock<Database> = OnceLock::new(); static MONGO_DATABASE: OnceCell<Database> = OnceCell::const_new();
pub async fn init() { pub async fn init() {
let mut options = let mut options =

1
swordfish/Cargo.toml
View File

@ -9,7 +9,6 @@ edition = "2021"
dotenvy = "0.15.7" dotenvy = "0.15.7"
image = "0.24.7" image = "0.24.7"
leptess = "0.14.0" leptess = "0.14.0"
regex = "1.10.2"
rusty-tesseract = "1.1.9" rusty-tesseract = "1.1.9"
serde = "1.0.193" serde = "1.0.193"
serenity = { version = "0.12.0", features = ["builder"] } serenity = { version = "0.12.0", features = ["builder"] }

106
swordfish/src/katana.rs
View File

@ -3,19 +3,17 @@ use crate::tesseract::{libtesseract, subprocess};
use crate::CONFIG; use crate::CONFIG;
use image::imageops::colorops::contrast_in_place; use image::imageops::colorops::contrast_in_place;
use image::io::Reader as ImageReader; use image::io::Reader as ImageReader;
use image::{DynamicImage, ImageFormat}; use image::{DynamicImage, GenericImage, GenericImageView, ImageBuffer, ImageFormat, Rgba};
use regex::Regex;
use serenity::all::Context; use serenity::all::Context;
use serenity::model::channel::Message; use serenity::model::channel::Message;
use std::io::Cursor; use std::io::Cursor;
use std::sync::LazyLock;
use swordfish_common::database::katana as db; use swordfish_common::database::katana as db;
use swordfish_common::structs::{Character, DroppedCard}; use swordfish_common::structs::{Character, DroppedCard};
use swordfish_common::{error, trace, warn}; use swordfish_common::{error, trace, warn};
use tokio::task; use tokio::task;
use tokio::time::Instant; use tokio::time::Instant;
static ALLOWED_CHARS_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"[\-!.':() ]").unwrap()); const ALLOWED_CHARS: [char; 10] = [' ', '-', '.', '!', ':', '(', ')', '\'', '/', '\''];
const CARD_NAME_X_OFFSET: u32 = 22; const CARD_NAME_X_OFFSET: u32 = 22;
const CARD_NAME_Y_OFFSET: u32 = 28; const CARD_NAME_Y_OFFSET: u32 = 28;
const CARD_NAME_WIDTH: u32 = 202 - CARD_NAME_X_OFFSET; const CARD_NAME_WIDTH: u32 = 202 - CARD_NAME_X_OFFSET;
@ -45,6 +43,7 @@ fn fix_tesseract_string(text: &mut String) {
// e.g. "We Never Learn\nN" -> "We Never Learn" // e.g. "We Never Learn\nN" -> "We Never Learn"
trace!("Text: {}", text); trace!("Text: {}", text);
if text.ends_with("\nN") { if text.ends_with("\nN") {
text.truncate(text.len() - 2);
for _ in 0..2 { for _ in 0..2 {
text.pop(); text.pop();
} }
@ -56,11 +55,10 @@ fn fix_tesseract_string(text: &mut String) {
// Workaround for a bug the text // Workaround for a bug the text
trace!("Text: {}", text); trace!("Text: {}", text);
if text.starts_with("- ") || text.starts_with("-.") { if text.starts_with("- ") || text.starts_with("-.") {
text.remove(0); text.drain(0..2);
text.remove(0);
} }
// Remove the first character if it is not alphanumeric // Remove the first character if it is not alphanumeric
if !text.clone().chars().nth(0).unwrap().is_ascii_alphanumeric() { if !text.starts_with(|c: char| c.is_ascii_alphanumeric()) {
text.remove(0); text.remove(0);
} }
// Workaround IR -> Ik // Workaround IR -> Ik
@ -81,11 +79,13 @@ fn fix_tesseract_string(text: &mut String) {
} }
} }
// Workaround for "\n." (and others in the future) // Workaround for "\n." (and others in the future)
for (i, c) in text.clone().chars().enumerate() { let text_clone = text.clone();
let mut clone_chars = text_clone.chars();
for (i, c) in clone_chars.clone().enumerate() {
if c != '\n' { if c != '\n' {
continue; continue;
} }
let prev_char = match text.chars().nth(i - 1) { let prev_char = match clone_chars.nth(i - 1) {
Some(c) => c, Some(c) => c,
None => continue, None => continue,
}; };
@ -97,7 +97,7 @@ fn fix_tesseract_string(text: &mut String) {
} }
// Fix for "Asobi ni Iku lo Asobi ni Oide" -> "Asobi ni Iku yo! Asobi ni Oide" // Fix for "Asobi ni Iku lo Asobi ni Oide" -> "Asobi ni Iku yo! Asobi ni Oide"
if prev_char == 'l' { if prev_char == 'l' {
let prev_prev_char = match text.chars().nth(i - 2) { let prev_prev_char = match clone_chars.nth(i - 2) {
Some(c) => c, Some(c) => c,
None => continue, None => continue,
}; };
@ -109,7 +109,7 @@ fn fix_tesseract_string(text: &mut String) {
text.insert_str(i - 2, "yo!") text.insert_str(i - 2, "yo!")
} }
} }
let next_char = match text.chars().nth(i + 1) { let next_char = match clone_chars.nth(i + 1) {
Some(c) => c, Some(c) => c,
None => break, None => break,
}; };
@ -125,7 +125,7 @@ fn fix_tesseract_string(text: &mut String) {
} }
// Remove all non-alphanumeric characters // Remove all non-alphanumeric characters
trace!("Text: {}", text); trace!("Text: {}", text);
text.retain(|c| ALLOWED_CHARS_REGEX.is_match(&c.to_string()) || c.is_ascii_alphanumeric()); text.retain(|c| ALLOWED_CHARS.contains(&c) || c.is_ascii_alphanumeric());
// Fix "mn" -> "III" // Fix "mn" -> "III"
trace!("Text: {}", text); trace!("Text: {}", text);
if text.ends_with("mn") { if text.ends_with("mn") {
@ -157,20 +157,29 @@ fn fix_tesseract_string(text: &mut String) {
} }
// Workaround if the first character is a space // Workaround if the first character is a space
trace!("Text: {}", text); trace!("Text: {}", text);
while text.starts_with(" ") { let mut leading_spaces = 0;
let mut ending_spaces = 0;
while text.starts_with(|c: char| c.is_whitespace()) {
trace!("Removing leading space"); trace!("Removing leading space");
text.remove(0); leading_spaces += 1;
} }
// Workaround if the last character is a space // Workaround if the last character is a space
trace!("Text: {}", text); while text.ends_with(|c: char| c.is_whitespace()) {
while text.ends_with(" ") {
trace!("Removing ending space"); trace!("Removing ending space");
text.pop(); ending_spaces += 1;
}
// Remove the ending spaces
if ending_spaces > 0 {
text.truncate(text.len() - ending_spaces);
}
// Remove the leading spaces
if leading_spaces > 0 {
text.drain(0..leading_spaces);
} }
trace!("Text (final): {}", text); trace!("Text (final): {}", text);
} }
fn save_image_if_trace(img: &image::DynamicImage, path: &str) { fn save_image_if_trace(img: &DynamicImage, path: &str) {
let log_lvl = CONFIG.get().unwrap().log.level.as_str(); let log_lvl = CONFIG.get().unwrap().log.level.as_str();
if log_lvl == "trace" { if log_lvl == "trace" {
match img.save(path) { match img.save(path) {
@ -184,21 +193,47 @@ fn save_image_if_trace(img: &image::DynamicImage, path: &str) {
} }
} }
fn image_with_white_padding(im: DynamicImage) -> DynamicImage {
// Partially copied from https://github.com/PureSci/nori/blob/main/rust-workers/src/drop.rs#L102C1-L121C6
let mut new_im: DynamicImage =
ImageBuffer::<Rgba<u8>, Vec<u8>>::new(im.width() + 14, im.height() + 14).into();
let white = Rgba([255, 255, 255, 255]);
for y in 0..im.height() {
for x in 0..im.width() {
let p = im.get_pixel(x, y);
new_im.put_pixel(x + 7, y + 7, p.to_owned());
}
}
for y in 0..7 {
for x in 0..im.width() + 14 {
new_im.put_pixel(x, y, white);
new_im.put_pixel(x, y + im.height() + 7, white);
}
}
for x in 0..7 {
for y in 7..im.height() + 7 {
new_im.put_pixel(x, y, white);
new_im.put_pixel(x + im.width() + 7, y, white);
}
}
new_im
}
pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) -> DroppedCard { pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) -> DroppedCard {
trace!("Spawning threads for analyzing card..."); trace!("Spawning threads for analyzing card...");
// Read the name and the series // Read the name and the series
let card_clone = card.clone(); let card_clone = card.clone();
let name_thread = task::spawn_blocking(move || unsafe { let name_thread = task::spawn_blocking(move || {
// let mut leptess = // let mut leptess =
// libtesseract::init_tesseract(false).expect("Failed to initialize Tesseract"); // libtesseract::init_tesseract(false).expect("Failed to initialize Tesseract");
let binding = libtesseract::get_tesseract(); let binding = unsafe { libtesseract::get_tesseract() };
let mut leptess = binding.lock().unwrap(); let mut leptess = binding.lock().unwrap();
let name_img = card_clone.crop_imm( let name_img = image_with_white_padding(card_clone.crop_imm(
CARD_NAME_X_OFFSET, CARD_NAME_X_OFFSET,
CARD_NAME_Y_OFFSET, CARD_NAME_Y_OFFSET,
CARD_NAME_WIDTH, CARD_NAME_WIDTH,
CARD_NAME_HEIGHT, CARD_NAME_HEIGHT,
); ));
let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new()); let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
match name_img.write_to(&mut buffer, ImageFormat::Png) { match name_img.write_to(&mut buffer, ImageFormat::Png) {
Ok(_) => {} Ok(_) => {}
@ -206,24 +241,27 @@ pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) ->
panic!("{}", format!("Failed to write image: {:?}", why)); panic!("{}", format!("Failed to write image: {:?}", why));
} }
}; };
save_image_if_trace(&name_img, format!("debug/4-{}-name.png", count).as_str()); save_image_if_trace(
&name_img,
format!("debug/4-libtesseract-{}-name.png", count).as_str(),
);
leptess.set_image_from_mem(&buffer.get_mut()).unwrap(); leptess.set_image_from_mem(&buffer.get_mut()).unwrap();
let mut name_str = leptess.get_utf8_text().expect("Failed to read name"); let mut name_str = leptess.get_utf8_text().expect("Failed to read name");
fix_tesseract_string(&mut name_str); fix_tesseract_string(&mut name_str);
name_str name_str
}); });
let card_clone = card.clone(); let card_clone = card.clone();
let series_thread = task::spawn_blocking(move || unsafe { let series_thread = task::spawn_blocking(move || {
// let mut leptess = // let mut leptess =
// libtesseract::init_tesseract(false).expect("Failed to initialize Tesseract"); // libtesseract::init_tesseract(false).expect("Failed to initialize Tesseract");
let binding = libtesseract::get_tesseract(); let binding = unsafe { libtesseract::get_tesseract() };
let mut leptess = binding.lock().unwrap(); let mut leptess = binding.lock().unwrap();
let series_img = card_clone.crop_imm( let series_img = image_with_white_padding(card_clone.crop_imm(
CARD_SERIES_X_OFFSET, CARD_SERIES_X_OFFSET,
CARD_SERIES_Y_OFFSET, CARD_SERIES_Y_OFFSET,
CARD_SERIES_WIDTH, CARD_SERIES_WIDTH,
CARD_SERIES_HEIGHT, CARD_SERIES_HEIGHT,
); ));
let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new()); let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
match series_img.write_to(&mut buffer, ImageFormat::Png) { match series_img.write_to(&mut buffer, ImageFormat::Png) {
Ok(_) => {} Ok(_) => {}
@ -233,10 +271,10 @@ pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) ->
}; };
save_image_if_trace( save_image_if_trace(
&series_img, &series_img,
format!("debug/4-{}-series.png", count).as_str(), format!("debug/4-libtesseract-{}-series.png", count).as_str(),
); );
leptess.set_image_from_mem(&buffer.get_mut()).unwrap(); leptess.set_image_from_mem(&buffer.get_mut()).unwrap();
let mut series_str = leptess.get_utf8_text().expect("Failed to read name"); let mut series_str = leptess.get_utf8_text().expect("Failed to read series");
fix_tesseract_string(&mut series_str); fix_tesseract_string(&mut series_str);
series_str series_str
}); });
@ -275,12 +313,12 @@ pub async fn analyze_card_subprocess(card: image::DynamicImage, count: u32) -> D
// Read the name and the series // Read the name and the series
let card_clone = card.clone(); let card_clone = card.clone();
let name_thread = task::spawn_blocking(move || { let name_thread = task::spawn_blocking(move || {
let name_img = card_clone.crop_imm( let name_img = image_with_white_padding(card_clone.crop_imm(
CARD_NAME_X_OFFSET, CARD_NAME_X_OFFSET,
CARD_NAME_Y_OFFSET, CARD_NAME_Y_OFFSET,
CARD_NAME_WIDTH, CARD_NAME_WIDTH,
CARD_NAME_HEIGHT, CARD_NAME_HEIGHT,
); ));
let img = subprocess::Image::from_dynamic_image(&name_img).unwrap(); let img = subprocess::Image::from_dynamic_image(&name_img).unwrap();
save_image_if_trace( save_image_if_trace(
&name_img, &name_img,
@ -292,12 +330,12 @@ pub async fn analyze_card_subprocess(card: image::DynamicImage, count: u32) -> D
}); });
let card_clone = card.clone(); let card_clone = card.clone();
let series_thread = task::spawn_blocking(move || { let series_thread = task::spawn_blocking(move || {
let series_img = card_clone.crop_imm( let series_img = image_with_white_padding(card_clone.crop_imm(
CARD_SERIES_X_OFFSET, CARD_SERIES_X_OFFSET,
CARD_SERIES_Y_OFFSET, CARD_SERIES_Y_OFFSET,
CARD_SERIES_WIDTH, CARD_SERIES_WIDTH,
CARD_SERIES_HEIGHT, CARD_SERIES_HEIGHT,
); ));
let img = subprocess::Image::from_dynamic_image(&series_img).unwrap(); let img = subprocess::Image::from_dynamic_image(&series_img).unwrap();
save_image_if_trace( save_image_if_trace(
&series_img, &series_img,
@ -368,7 +406,7 @@ pub async fn analyze_drop_message(message: &Message) -> Result<Vec<DroppedCard>,
img = img.grayscale(); img = img.grayscale();
save_image_if_trace(&img, "debug/1-grayscale.png"); save_image_if_trace(&img, "debug/1-grayscale.png");
trace!("Increasing contrast of the image..."); trace!("Increasing contrast of the image...");
contrast_in_place(&mut img, 127.0); contrast_in_place(&mut img, 127.0 / 4.0);
save_image_if_trace(&img, "debug/2-contrast.png"); save_image_if_trace(&img, "debug/2-contrast.png");
// Cropping cards // Cropping cards
let distance = 257 - 29 + 305 - 259; let distance = 257 - 29 + 305 - 259;

4
swordfish/src/main.rs
View File

@ -8,9 +8,9 @@ use serenity::model::channel::Message;
use serenity::prelude::*; use serenity::prelude::*;
use std::env; use std::env;
use std::path::Path; use std::path::Path;
use std::sync::OnceLock;
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
use swordfish_common::*; use swordfish_common::*;
use tokio::sync::OnceCell;
use crate::config::Config; use crate::config::Config;
use crate::tesseract::libtesseract; use crate::tesseract::libtesseract;
@ -23,7 +23,7 @@ mod template;
mod tesseract; mod tesseract;
const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish"; const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish";
static CONFIG: OnceLock<Config> = OnceLock::new(); static CONFIG: OnceCell<Config> = OnceCell::const_new();
#[group] #[group]
#[commands(ping, debug, info)] #[commands(ping, debug, info)]