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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

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

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

@ -1,12 +1,12 @@
use crate::database;
use crate::structs::Character;
use mongodb::Collection;
use std::sync::OnceLock;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::OnceCell;
use tokio::task;
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

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

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

1
swordfish/Cargo.toml
View File

@ -9,7 +9,6 @@ edition = "2021"
dotenvy = "0.15.7"
image = "0.24.7"
leptess = "0.14.0"
regex = "1.10.2"
rusty-tesseract = "1.1.9"
serde = "1.0.193"
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 image::imageops::colorops::contrast_in_place;
use image::io::Reader as ImageReader;
use image::{DynamicImage, ImageFormat};
use regex::Regex;
use image::{DynamicImage, GenericImage, GenericImageView, ImageBuffer, ImageFormat, Rgba};
use serenity::all::Context;
use serenity::model::channel::Message;
use std::io::Cursor;
use std::sync::LazyLock;
use swordfish_common::database::katana as db;
use swordfish_common::structs::{Character, DroppedCard};
use swordfish_common::{error, trace, warn};
use tokio::task;
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_Y_OFFSET: u32 = 28;
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"
trace!("Text: {}", text);
if text.ends_with("\nN") {
text.truncate(text.len() - 2);
for _ in 0..2 {
text.pop();
}
@ -56,11 +55,10 @@ fn fix_tesseract_string(text: &mut String) {
// Workaround for a bug the text
trace!("Text: {}", text);
if text.starts_with("- ") || text.starts_with("-.") {
text.remove(0);
text.remove(0);
text.drain(0..2);
}
// 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);
}
// Workaround IR -> Ik
@ -81,11 +79,13 @@ fn fix_tesseract_string(text: &mut String) {
}
}
// 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' {
continue;
}
let prev_char = match text.chars().nth(i - 1) {
let prev_char = match clone_chars.nth(i - 1) {
Some(c) => c,
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"
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,
None => continue,
};
@ -109,7 +109,7 @@ fn fix_tesseract_string(text: &mut String) {
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,
None => break,
};
@ -125,7 +125,7 @@ fn fix_tesseract_string(text: &mut String) {
}
// Remove all non-alphanumeric characters
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"
trace!("Text: {}", text);
if text.ends_with("mn") {
@ -157,20 +157,29 @@ fn fix_tesseract_string(text: &mut String) {
}
// Workaround if the first character is a space
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");
text.remove(0);
leading_spaces += 1;
}
// Workaround if the last character is a space
trace!("Text: {}", text);
while text.ends_with(" ") {
while text.ends_with(|c: char| c.is_whitespace()) {
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);
}
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();
if log_lvl == "trace" {
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 {
trace!("Spawning threads for analyzing card...");
// Read the name and the series
let card_clone = card.clone();
let name_thread = task::spawn_blocking(move || unsafe {
let name_thread = task::spawn_blocking(move || {
// let mut leptess =
// 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 name_img = card_clone.crop_imm(
let name_img = image_with_white_padding(card_clone.crop_imm(
CARD_NAME_X_OFFSET,
CARD_NAME_Y_OFFSET,
CARD_NAME_WIDTH,
CARD_NAME_HEIGHT,
);
));
let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
match name_img.write_to(&mut buffer, ImageFormat::Png) {
Ok(_) => {}
@ -206,24 +241,27 @@ pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) ->
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();
let mut name_str = leptess.get_utf8_text().expect("Failed to read name");
fix_tesseract_string(&mut name_str);
name_str
});
let card_clone = card.clone();
let series_thread = task::spawn_blocking(move || unsafe {
let series_thread = task::spawn_blocking(move || {
// let mut leptess =
// 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 series_img = card_clone.crop_imm(
let series_img = image_with_white_padding(card_clone.crop_imm(
CARD_SERIES_X_OFFSET,
CARD_SERIES_Y_OFFSET,
CARD_SERIES_WIDTH,
CARD_SERIES_HEIGHT,
);
));
let mut buffer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
match series_img.write_to(&mut buffer, ImageFormat::Png) {
Ok(_) => {}
@ -233,10 +271,10 @@ pub async fn analyze_card_libtesseract(card: image::DynamicImage, count: u32) ->
};
save_image_if_trace(
&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();
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);
series_str
});
@ -275,12 +313,12 @@ pub async fn analyze_card_subprocess(card: image::DynamicImage, count: u32) -> D
// Read the name and the series
let card_clone = card.clone();
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_Y_OFFSET,
CARD_NAME_WIDTH,
CARD_NAME_HEIGHT,
);
));
let img = subprocess::Image::from_dynamic_image(&name_img).unwrap();
save_image_if_trace(
&name_img,
@ -292,12 +330,12 @@ pub async fn analyze_card_subprocess(card: image::DynamicImage, count: u32) -> D
});
let card_clone = card.clone();
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_Y_OFFSET,
CARD_SERIES_WIDTH,
CARD_SERIES_HEIGHT,
);
));
let img = subprocess::Image::from_dynamic_image(&series_img).unwrap();
save_image_if_trace(
&series_img,
@ -368,7 +406,7 @@ pub async fn analyze_drop_message(message: &Message) -> Result<Vec<DroppedCard>,
img = img.grayscale();
save_image_if_trace(&img, "debug/1-grayscale.png");
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");
// Cropping cards
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 std::env;
use std::path::Path;
use std::sync::OnceLock;
use std::time::{SystemTime, UNIX_EPOCH};
use swordfish_common::*;
use tokio::sync::OnceCell;
use crate::config::Config;
use crate::tesseract::libtesseract;
@ -23,7 +23,7 @@ mod template;
mod tesseract;
const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish";
static CONFIG: OnceLock<Config> = OnceLock::new();
static CONFIG: OnceCell<Config> = OnceCell::const_new();
#[group]
#[commands(ping, debug, info)]