seoxi/watcat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved color processing, improved logging

Browse Source
This commit is contained in:
Seoxi Ryouko
2025-05-16 10:50:00 -05:00
parent f66f89dc04
commit 88d4b19655
3 changed files with 175 additions and 708 deletions
Download Patch File Download Diff File Expand all files Collapse all files

743
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

4
Cargo.toml
View File

@ -10,6 +10,7 @@ dotenvy = "0.15.7"
reqwest = "0.12"
reqwest-middleware = "0.4.2"
http-cache-reqwest = "0.15.1"
palette = "0.7.6"
[dependencies.magick_rust]
path = "./magick-rust"
@ -21,6 +22,3 @@ features = ["macros", "rt-multi-thread"]
[dependencies.sqlx]
version = "0.8.5"
features = ["sqlite", "runtime-tokio"]
[dependencies.imagetext]
git = "https://github.com/nathanielfernandes/imagetext"

136
src/main.rs
View File

@ -16,6 +16,10 @@ use magick_rust::{ColorspaceType, CompositeOperator, MagickWand, PixelWand, magi
use std::env;
use std::sync::Once;
use palette::color_difference::ImprovedCiede2000;
use palette::lab::Lab;
use palette::white_point::D65;
extern crate sqlx;
use sqlx::sqlite::SqlitePool;
@ -28,12 +32,34 @@ enum Reply {
Text(String),
}
macro_rules! now {
() => {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs_f32() as usize
};
}
macro_rules! log {
($a:expr) => {
if env::var("DEBUG").unwrap_or("0".to_string()) == "1" {
println!(concat!("debug: {} ", $a), now!())
}
};
($a:expr, $($b:expr),+) => {
if env::var("DEBUG").unwrap_or("0".to_string()) == "1" {
println!(concat!("debug: {} ", $a), now!(), $($b),+)
}
};
}
macro_rules! handle_magick_result {
($a: expr, $b: literal) => {
match $a {
Ok(_) => {
#[cfg(debug_assertions)]
println!("debug: {} run successfully", stringify!($a));
log!("{} run successfully", stringify!($a));
}
Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)),
}
@ -49,6 +75,40 @@ macro_rules! handle_magick_option {
};
}
// this function is almost assuredly the least idiomatic rust code I will ever write.
// but until I know what the "correct" way is, it'll have to do
fn parse_cielab(s: String) -> Result<(f32, f32, f32), &'static str> {
if s.len() < 13 || &s[0..7] != "cielab(" || s.chars().nth(s.len() - 1).unwrap() != ')' {
return Err("Error parsing cielab string");
}
let vals = s[7..(s.len() - 1)]
.split(",")
.map(|x| match x.parse() {
Ok(f) => f,
Err(_) => f32::NAN,
})
.collect::<Vec<f32>>();
if vals.iter().any(|x| x.is_nan()) {
return Err("Error parsing cielab string");
}
match &vals[..] {
&[a, b, c] => Ok((a, b, c)),
_ => Err("Error parsing cielab string"),
}
}
fn validate_color(col: &PixelWand) -> Option<Lab> {
let color_str = match col.get_color_as_string() {
Ok(s) => s,
Err(_) => return None,
};
let color_raw = match parse_cielab(color_str) {
Ok(f) => f,
Err(_) => return None,
};
Some(Lab::<D65>::from_components(color_raw))
}
async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Reply {
let lastfm_user = match arg {
"" => get_lastfm_username(ctx, id).await,
@ -110,7 +170,7 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
"Failed to set cluster colorspace"
);
handle_magick_result!(
art_wand_cluster.kmeans(5, 200, 0.001),
art_wand_cluster.kmeans(10, 200, 0.001),
"Failed to run kmeans"
);
let mut colors = handle_magick_option!(
@ -118,10 +178,29 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
"Failed to get color histogram"
);
colors.sort_by_cached_key(|color| -(color.get_color_count() as isize));
let accent_color = match colors.len() {
1 => &colors[0],
_ => &colors[1],
};
let base_lab = validate_color(&colors[0]).expect("Failed to validate base color");
let other_colors = colors
.iter()
.filter_map(|col| {
let testing_color = validate_color(col)?;
match testing_color.improved_difference(base_lab) {
10.0.. => Some(testing_color),
_ => None,
}
})
.collect::<Vec<_>>();
let mut accent_color = PixelWand::new();
handle_magick_result!(
accent_color.set_color("#ffffff"),
"Failed to init accent color"
);
if let Some(color) = other_colors.first() {
let col = format!("cielab{:?}", color.into_components());
handle_magick_result!(
accent_color.set_color(col.as_str()),
"Failed to set accent color"
);
}
handle_magick_result!(white.set_color("#ffffff"), "Failed to init white");
handle_magick_result!(
@ -129,7 +208,7 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
"Failed to load mask"
);
handle_magick_result!(
mask_wand.opaque_paint_image(&white, accent_color, 0.0, false),
mask_wand.opaque_paint_image(&white, &accent_color, 0.0, false),
"Failed to paint accent mask"
);
@ -153,10 +232,6 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
"Failed to load dummy avatar"
),
}
handle_magick_result!(
avatar_wand.transform_image_colorspace(ColorspaceType::Lab),
"Failed to set avatar colorspace"
);
let avatar_mask_wand = MagickWand::new();
handle_magick_result!(
avatar_mask_wand.read_image_blob(include_bytes!("avatar_mask.png")),
@ -170,6 +245,10 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
avatar_wand.adaptive_resize_image(64, 64),
"Failed to resize avatar"
);
handle_magick_result!(
avatar_wand.transform_image_colorspace(ColorspaceType::Lab),
"Failed to set avatar colorspace"
);
handle_magick_result!(
main_wand.new_image(548, 147, &colors[0]),
@ -195,6 +274,16 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
main_wand.compose_images(&mask_wand, CompositeOperator::SrcOver, false, 401, 0),
"Failed to combine mask"
);
handle_magick_result!(
main_wand.compose_images(
&art_wand_cluster,
CompositeOperator::SrcOver,
false,
160,
12
),
"Failed to combine debug image"
);
handle_magick_result!(
main_wand.compose_images(&avatar_wand, CompositeOperator::SrcOver, false, 473, 73),
"Failed to combine avatar image"
@ -214,8 +303,14 @@ impl EventHandler for Handler {
let cmd = cmd_iter.next().unwrap_or("");
let arg = cmd_iter.next().unwrap_or("");
let resp = match cmd {
".fmi" => Some(fmi(&ctx, arg, msg.author.id, msg.author.avatar_url()).await),
".set" => Some(set(&ctx, arg, msg.author.id).await),
".fmi" => {
log!("{} received", msg.content);
Some(fmi(&ctx, arg, msg.author.id, msg.author.avatar_url()).await)
}
".set" => {
log!("{} received", msg.content);
Some(set(&ctx, arg, msg.author.id).await)
}
_ => None,
};
if let Some(reply) = resp {
@ -250,6 +345,7 @@ struct DBResponse {
}
async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
log!("get db user {}", id);
let data = ctx.data.write().await;
let pool = data
.get::<PoolContainer>()
@ -267,12 +363,16 @@ async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
.fetch_one(pool)
.await;
match resp {
Ok(r) => Some(r.lastfm_username),
Ok(r) => {
log!("got user {}", r.lastfm_username);
Some(r.lastfm_username)
}
Err(_) => None,
}
}
async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
log!("set db user {} {}", id, user);
let data = ctx.data.write().await;
let pool = data
.get::<PoolContainer>()
@ -292,16 +392,14 @@ async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
}
async fn get_image(ctx: &Context, url: &str) -> Result<Vec<u8>, reqwest_middleware::Error> {
#[cfg(debug_assertions)]
println!("debug: get {}", url);
log!("get {}", url);
let data = ctx.data.write().await;
let http = data
.get::<HttpContainer>()
.expect("Failed to get http client");
match http.get(url).send().await {
Ok(resp) => {
#[cfg(debug_assertions)]
println!("debug: response received");
log!("response received");
Ok(resp
.bytes()
.await