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

Added color cache

Browse Source
This commit is contained in:
Seoxi Ryouko
2025-07-31 04:10:02 -05:00
parent 42270b18f0
commit d3ad04bc33
5 changed files with 121 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
.env.example
View File

@ -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=

2
Cargo.lock generated
View File

@ -4270,7 +4270,7 @@ dependencies = [
[[package]]
name = "watcat"
version = "0.5.0"
version = "0.6.0"
dependencies = [
"arabic_reshaper",
"dotenvy",

2
Cargo.toml
View File

@ -1,6 +1,6 @@
[package]
name = "watcat"
version = "0.5.0"
version = "0.6.0"
edition = "2024"
[dependencies]

2
README.md
View File

@ -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?

182
src/main.rs
View File

@ -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<String>) -> 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<String>) -> 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::<ColorsHaver>().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::<Vec<_>>();
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::<Vec<_>>();
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::<ColorsHaver>().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::<D65>::new(0.0, 0.0, 0.0).improved_difference(color);
let white_delta_e = Lab::<D65>::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<String>) -> 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<String>) -> 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<Arc<Mutex<ClientCredsSpotify>>>;
}
struct ColorCache {
map: Arc<Mutex<HashMap<String, (String, String)>>>,
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<String> {
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::<PoolHaver>()
.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<Vec<u8>, 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::<HttpHaver>().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::<usize>().unwrap_or(1000000),
};
let mut discord_client = Client::builder(&token, intents)
.event_handler(Handler)
.await
@ -663,6 +702,7 @@ async fn main() {
data.insert::<HttpHaver>(http);
data.insert::<FontsHaver>((regular_fonts, bold_fonts));
data.insert::<SpotifyHaver>(spotify.map(Mutex::new).map(Arc::new));
data.insert::<ColorsHaver>(colors_cache);
}
if let Err(why) = discord_client.start().await {