From d3ad04bc33574bb9572ede301163681bf9174709 Mon Sep 17 00:00:00 2001 From: Seoxi Ryouko Date: Thu, 31 Jul 2025 04:10:02 -0500 Subject: [PATCH] Added color cache --- .env.example | 7 ++ Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 2 +- src/main.rs | 182 +++++++++++++++++++++++++++++++-------------------- 5 files changed, 121 insertions(+), 74 deletions(-) diff --git a/.env.example b/.env.example index ae57049..2378445 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,10 @@ DISCORD_TOKEN= LASTFM_API_KEY= DATABASE_URL=sqlite:watcat.db + +# optional, for spotify support +RSPOTIFY_CLIENT_ID= +RSPOTIFY_CLIENT_SECRET= + +# optional, to set cache size, set to 0 to disable +MAX_CACHE_SIZE= diff --git a/Cargo.lock b/Cargo.lock index 5aafd89..86cca8c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4270,7 +4270,7 @@ dependencies = [ [[package]] name = "watcat" -version = "0.5.0" +version = "0.6.0" dependencies = [ "arabic_reshaper", "dotenvy", diff --git a/Cargo.toml b/Cargo.toml index bcf9528..3f821d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "watcat" -version = "0.5.0" +version = "0.6.0" edition = "2024" [dependencies] diff --git a/README.md b/README.md index ad7f51b..48fce84 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ TODO + ~~refactor to combine repeated parts of art, url, and fmi.~~ + use macro for commands + use Option instead of empty string sentinel value -+ add cache for album art colors ++ ~~add cache for album art colors~~ + ~~add spotify link fetching~~ + make various parameters user-configurable + use tiny-skia instead of magick-rust? diff --git a/src/main.rs b/src/main.rs index 79e8924..1613aae 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ use magick_rust::{ }; use std::env; +use std::collections::HashMap; use std::sync::{Arc, Once}; use tokio::sync::Mutex; @@ -69,8 +70,9 @@ macro_rules! log { macro_rules! handle_magick_result { ($a: expr, $b: literal) => { match $a { - Ok(_) => { + Ok(a) => { log!("{} run successfully", stringify!($a)); + a } Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)), } @@ -179,7 +181,6 @@ async fn spot(ctx: &Context, arg: &str, id: UserId) -> Reply { return Reply::Text("Unable to use Spotify command: Contact bot Administrator.".to_owned()) } let spotify_arc_mutex = spotify_option.as_mut().unwrap(); - //let spotify_arc_mutex_ref = Arc::clone(&spotify_arc_mutex); let spotify_client = spotify_arc_mutex.lock().await; spotify_client.request_token().await.unwrap(); let search = spotify_client.search(format!("{} {}", track.name, track.artist.name).as_str(), SearchType::Track, None, None, None, None).await; @@ -209,11 +210,8 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option) -> Re None => return Reply::Text("Error: getting image uri failed".to_owned()), }; - let image = match get_image(ctx, image_uri.as_str()).await { - Ok(i) => i, - Err(e) => return Reply::Text(format!("{e}")), - }; let mut base_color = PixelWand::new(); + let mut accent_color = PixelWand::new(); let mut white = PixelWand::new(); let main_wand = MagickWand::new(); let art_wand = MagickWand::new(); @@ -221,67 +219,94 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option) -> Re let mask_wand = MagickWand::new(); let text_image_wand = MagickWand::new(); - handle_magick_result!( - base_color.set_color("#7f7f7f"), - "Failed to set init base color" - ); - handle_magick_result!( - art_wand.read_image_blob(image), - "Failed to read image from uri" - ); - handle_magick_result!( - art_wand.transform_image_colorspace(ColorspaceType::Lab), - "Failed to set art colorspace" - ); + let image = match get_image(ctx, image_uri.as_str()).await { + Ok(i) => i, + Err(e) => return Reply::Text(format!("{e}")), + }; + handle_magick_result!( + art_wand.read_image_blob(image), + "Failed to read image from uri" + ); + handle_magick_result!( + art_wand.transform_image_colorspace(ColorspaceType::Lab), + "Failed to set art colorspace" + ); + + handle_magick_result!( + base_color.set_color("#7f7f7f"), + "Failed to set init base color" + ); + + let color_cache_entry = { + let data = ctx.data.write().await; + let color_cache = data.get::().expect("Failed to have colors option"); + if color_cache.max > 0 { + let color_cache_map = color_cache.map.lock().await; + color_cache_map.get(image_uri.as_str()).map(|s| s.to_owned()) + } else { + None + } + }; + - let (art_width, art_height) = (art_wand.get_image_width(), art_wand.get_image_height()); - handle_magick_result!( - art_wand_cluster.new_image(art_width, art_height, &base_color), - "Failed to init cluster image" - ); - handle_magick_result!( - art_wand_cluster.compose_images(&art_wand, CompositeOperator::SrcOver, true, 0, 0), - "Failed to copy onto cluster image" - ); - // least stupid way I could figure out to do this, since it seems that MagickWand::new_from_wand() doesn't work - handle_magick_result!( - art_wand_cluster.set_image_colorspace(ColorspaceType::Lab), - "Failed to set cluster colorspace" - ); - handle_magick_result!( - art_wand_cluster.kmeans(10, 200, 0.001), - "Failed to run kmeans" - ); - let mut colors = handle_magick_option!( - art_wand_cluster.get_image_histogram(), - "Failed to get color histogram" - ); - colors.sort_by_cached_key(|color| -(color.get_color_count() as isize)); - 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::>(); - let mut accent_color = PixelWand::new(); - handle_magick_result!( - accent_color.set_color("cielab(0.0,0.0,0.0)"), - "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" - ); - } + let (base_lab, accent_lab) = match color_cache_entry { + Some(value) => value.to_owned(), + None => { + let (art_width, art_height) = (art_wand.get_image_width(), art_wand.get_image_height()); + handle_magick_result!( + art_wand_cluster.new_image(art_width, art_height, &base_color), + "Failed to init cluster image" + ); + handle_magick_result!( + art_wand_cluster.compose_images(&art_wand, CompositeOperator::SrcOver, true, 0, 0), + "Failed to copy onto cluster image" + ); + // least stupid way I could figure out to do this, since it seems that MagickWand::new_from_wand() doesn't work + handle_magick_result!( + art_wand_cluster.set_image_colorspace(ColorspaceType::Lab), + "Failed to set cluster colorspace" + ); + handle_magick_result!( + art_wand_cluster.kmeans(10, 200, 0.001), + "Failed to run kmeans" + ); + let mut colors = handle_magick_option!( + art_wand_cluster.get_image_histogram(), + "Failed to get color histogram" + ); + colors.sort_by_cached_key(|color| -(color.get_color_count() as isize)); + 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::>(); + let mut accent_lab = "cielab(0.0,0.0,0.0)".to_owned(); + if let Some(color) = other_colors.first() { + accent_lab = format!("cielab{:?}", color.into_components()); + } + let base_lab = handle_magick_result!(colors[0].get_color_as_string(), "Failed to serialize base color"); + let data = ctx.data.write().await; + let color_cache = data.get::().expect("Failed to have colors option"); + let mut color_cache_map = color_cache.map.lock().await; + if color_cache_map.keys().len() > color_cache.max { + let key: String = color_cache_map.iter().next().unwrap().0.to_owned(); + color_cache_map.remove(&key); + } + color_cache_map.insert(image_uri, (base_lab.clone(), accent_lab.clone())); + (base_lab, accent_lab) + } + }; - let use_dark_color = match validate_color(&colors[0]) { + handle_magick_result!(base_color.set_color(base_lab.as_str()), "Failed to deserialize base color"); + handle_magick_result!(accent_color.set_color(accent_lab.as_str()), "Failed to deserialize accent color"); + + let use_dark_color = match validate_color(&base_color) { Some(color) => { let black_delta_e = Lab::::new(0.0, 0.0, 0.0).improved_difference(color); let white_delta_e = Lab::::new(100.0, 0.0, 0.0).improved_difference(color); @@ -362,7 +387,7 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option) -> Re ); handle_magick_result!( - main_wand.new_image(image_width as usize, image_height as usize, &colors[0]), + main_wand.new_image(image_width as usize, image_height as usize, &base_color), "Failed to create wand" ); handle_magick_result!( @@ -373,10 +398,10 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option) -> Re art_wand.adaptive_resize_image(124, 124), "Failed to resize art" ); - handle_magick_result!( + /*handle_magick_result!( art_wand_cluster.adaptive_resize_image(124, 124), "Failed to resize art_cluster" - ); + );*/ handle_magick_result!( main_wand.compose_images(&art_wand, CompositeOperator::SrcOver, false, 12, 12), "Failed to combine art image" @@ -519,6 +544,15 @@ impl TypeMapKey for SpotifyHaver { type Value = Option>>; } +struct ColorCache { + map: Arc>>, + max: usize, +} +struct ColorsHaver; +impl TypeMapKey for ColorsHaver { + type Value = ColorCache; +} + struct DBResponse { lastfm_username: String, } @@ -552,7 +586,7 @@ async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option { async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) { log!("set db user {} {}", id, user); - let data = ctx.data.write().await; + let data = ctx.data.read().await; let pool = data .get::() .expect("Failed to get pool container"); @@ -571,8 +605,8 @@ async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) { } async fn get_image(ctx: &Context, url: &str) -> Result, reqwest_middleware::Error> { - log!("get {}", url); - let data = ctx.data.write().await; + log!("get {}", url); + let data = ctx.data.read().await; let http = data.get::().expect("Failed to get http client"); match http.get(url).send().await { Ok(resp) => { @@ -652,6 +686,11 @@ async fn main() { let _ = &spotify_client.request_token().await.unwrap(); } + let colors_cache: ColorCache = ColorCache { + map: Arc::new(Mutex::new(HashMap::new())), + max: env::var("MAX_CACHE_SIZE").unwrap_or("".to_owned()).parse::().unwrap_or(1000000), + }; + let mut discord_client = Client::builder(&token, intents) .event_handler(Handler) .await @@ -663,6 +702,7 @@ async fn main() { data.insert::(http); data.insert::((regular_fonts, bold_fonts)); data.insert::(spotify.map(Mutex::new).map(Arc::new)); + data.insert::(colors_cache); } if let Err(why) = discord_client.start().await {