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

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