From b45739b2e05793bb9c64e86fe74f97531c9c473b Mon Sep 17 00:00:00 2001 From: Daniel Bulant Date: Wed, 2 Oct 2024 22:41:34 +0200 Subject: [PATCH] add vibrant color detection --- Cargo.lock | 21 +++- Cargo.toml | 6 +- src/main.rs | 1 + src/spotify.rs | 43 ++++++- src/vibrancy/README.md | 1 + src/vibrancy/mod.rs | 7 ++ src/vibrancy/palette.rs | 114 +++++++++++++++++++ src/vibrancy/settings.rs | 21 ++++ src/vibrancy/vibrant.rs | 238 +++++++++++++++++++++++++++++++++++++++ 9 files changed, 445 insertions(+), 7 deletions(-) create mode 100644 src/vibrancy/README.md create mode 100644 src/vibrancy/mod.rs create mode 100644 src/vibrancy/palette.rs create mode 100644 src/vibrancy/settings.rs create mode 100644 src/vibrancy/vibrant.rs diff --git a/Cargo.lock b/Cargo.lock index 2046ee2..daae363 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -276,9 +276,13 @@ dependencies = [ name = "barrs" version = "0.1.0" dependencies = [ + "color_quant", "cushy", + "hsl", "image", + "itertools 0.10.5", "mpris", + "palette", "plotters", "reqwest", "tokio", @@ -1330,6 +1334,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" +[[package]] +name = "hsl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "575fb7f1167f3b88ed825e90eb14918ac460461fdeaa3965c6a50951dee1c970" + [[package]] name = "http" version = "1.1.0" @@ -1542,6 +1552,15 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "187674a687eed5fe42285b40c6291f9a01517d415fad1c3cbc6a9f778af7fcd4" +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.12.1" @@ -2760,7 +2779,7 @@ dependencies = [ "built", "cfg-if", "interpolate_name", - "itertools", + "itertools 0.12.1", "libc", "libfuzzer-sys", "log", diff --git a/Cargo.toml b/Cargo.toml index 604e2dd..e361557 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,4 +9,8 @@ tokio = { version = "1.40.0", features = ["rt", "rt-multi-thread"] } plotters = { version = "0.3.7", default-features = false } image = { version = "0.25.0", features = ["png"] } mpris = "2.0.1" -reqwest = "0.12.8" \ No newline at end of file +reqwest = "0.12.8" +color_quant = "1.0" +hsl = "0.1.1" +itertools = "0.10.0" +palette = "0.7.3" \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 8317e75..19b3c10 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use cushy::{Open, PendingApp, Run, TokioRuntime}; mod spotify; +mod vibrancy; fn main() -> cushy::Result { let mut app = PendingApp::new(TokioRuntime::default()); diff --git a/src/spotify.rs b/src/spotify.rs index 04e4df9..058aa2d 100644 --- a/src/spotify.rs +++ b/src/spotify.rs @@ -1,10 +1,12 @@ use std::{sync::{Arc, Mutex}, thread, time::Duration}; use mpris::{LoopStatus, PlaybackStatus, PlayerFinder}; -use cushy::{kludgine::{AnyTexture, LazyTexture}, styles::components::WidgetBackground, value::{Destination, Dynamic, IntoReader, Source}, widget::MakeWidget, widgets::Image, Open, PendingApp}; +use cushy::{kludgine::{AnyTexture, LazyTexture}, styles::{components::WidgetBackground, Color}, value::{Destination, Dynamic, IntoReader, Source}, widget::MakeWidget, widgets::Image, Open, PendingApp}; +use palette::Srgb; use reqwest::Client; -use image; +use image::{self, Rgb}; use tokio::{runtime, task::JoinHandle}; +use crate::vibrancy::Vibrancy; #[derive(PartialEq)] struct PlayingTrack { @@ -22,7 +24,7 @@ struct PlayingTrack { pub fn spotify_controls() -> impl MakeWidget { let (progress, track) = get_track_dynamics(); - let texture = get_texture_dynamic(track.clone()); + let (texture, vibrancy) = get_texture_dynamic(track.clone()); track.map_each(|track| { if let Some(track) = track { @@ -36,10 +38,14 @@ pub fn spotify_controls() -> impl MakeWidget { } }) .to_label() + .centered() + .pad() .and( Image::new(texture) ) .into_rows() + .with(&WidgetBackground, vibrancy.map_each(|vib| vib.dark.unwrap_or(Color::BLACK).into())) + .pad() } fn get_empty_texture() -> AnyTexture { @@ -83,15 +89,27 @@ fn tokio_runtime() -> &'static runtime::Handle { }) } -fn get_texture_dynamic(track: Dynamic>) -> Dynamic { +#[derive(Debug, PartialEq, Eq, Default)] +pub struct ImageVibrancy { + primary: Option, + dark: Option, + light: Option, + muted: Option, + dark_muted: Option, + light_muted: Option, +} + +fn get_texture_dynamic(track: Dynamic>) -> (Dynamic, Dynamic) { let client = Client::new(); let texture = Dynamic::new(get_empty_texture()); + let vibrancy = Dynamic::new(ImageVibrancy::default()); let prev_request_join = Arc::new(Mutex::new(None::>)); track.for_each({ let texture = texture.clone(); + let vibrancy = vibrancy.clone(); move |track| { if let Some(track) = track { let mut prev_request_join = prev_request_join.lock().unwrap(); @@ -99,12 +117,22 @@ fn get_texture_dynamic(track: Dynamic>) -> Dynamic>) -> Dynamic) -> Color { + Color::new(rgb[0], rgb[1], rgb[2], 255) } /// This spawns a new thread to track the current playing track and its progress. diff --git a/src/vibrancy/README.md b/src/vibrancy/README.md new file mode 100644 index 0000000..b6552aa --- /dev/null +++ b/src/vibrancy/README.md @@ -0,0 +1 @@ +Adapted from https://github.com/killercup/vibrant-rs. MIT License \ No newline at end of file diff --git a/src/vibrancy/mod.rs b/src/vibrancy/mod.rs new file mode 100644 index 0000000..708865c --- /dev/null +++ b/src/vibrancy/mod.rs @@ -0,0 +1,7 @@ + +mod settings; +mod palette; +mod vibrant; + +pub use vibrant::Vibrancy; +pub use palette::Palette; \ No newline at end of file diff --git a/src/vibrancy/palette.rs b/src/vibrancy/palette.rs new file mode 100644 index 0000000..2c3437d --- /dev/null +++ b/src/vibrancy/palette.rs @@ -0,0 +1,114 @@ +use std::fmt; +use std::collections::BTreeMap; + +use itertools::Itertools; +use image::{GenericImage, Pixel, Rgb, Rgba}; +use color_quant::NeuQuant; + +/// Palette of colors. +#[derive(Debug, Hash, PartialEq, Eq, Default)] +pub struct Palette { + /// Palette of Colors represented in RGB + pub palette: Vec>, + /// A map of indices in the palette to a count of pixels in approximately that color in the + /// original image. + pub pixel_counts: BTreeMap, +} + +impl Palette { + /// Create a new palett from an image + /// + /// Color count and quality are given straight to [color_quant], values should be between + /// 8...512 and 1...30 respectively. (By the way: 10 is a good default quality.) + /// + /// [color_quant]: https://github.com/PistonDevelopers/color_quant + pub fn new(image: &G, color_count: usize, quality: i32) -> Palette + where P: Sized + Pixel, + G: Sized + GenericImage + { + let pixels: Vec> = image.pixels() + .map(|(_, _, pixel)| pixel.to_rgba()) + .collect(); + + let mut flat_pixels: Vec = Vec::with_capacity(pixels.len()); + for rgba in &pixels { + if is_boring_pixel(&rgba) { + continue; + } + + for subpixel in rgba.channels() { + flat_pixels.push(*subpixel); + } + } + + let quant = NeuQuant::new(quality, color_count, &flat_pixels); + + let pixel_counts = pixels.iter() + .map(|rgba| quant.index_of(&rgba.channels())) + .fold(BTreeMap::new(), + |mut acc, pixel| { + *acc.entry(pixel).or_insert(0) += 1; + acc + }); + + let palette: Vec> = quant.color_map_rgba() + .iter() + .chunks(4) + .into_iter() + .map(|rgba_iter| { + let rgba_slice: Vec = rgba_iter.cloned().collect(); + Rgba::from_slice(&rgba_slice).clone().to_rgb() + }) + .unique() + .collect(); + + Palette { + palette: palette, + pixel_counts: pixel_counts, + } + } + + fn frequency_of(&self, color: &Rgb) -> usize { + let index = self.palette.iter().position(|x| x.channels() == color.channels()); + if let Some(index) = index { + *self.pixel_counts.get(&index).unwrap_or(&0) + } else { + 0 + } + } + + /// Change ordering of colors in palette to be of frequency using the pixel count. + pub fn sort_by_frequency(&self) -> Self { + let mut colors = self.palette.clone(); + colors.sort_by(|a, b| self.frequency_of(&a).cmp(&self.frequency_of(&b))); + + Palette { + palette: colors, + pixel_counts: self.pixel_counts.clone(), + } + } +} + +fn is_boring_pixel(pixel: &Rgba) -> bool { + let (r, g, b, a) = (pixel[0], pixel[1], pixel[2], pixel[3]); + + // If pixel is mostly opaque and not white + const MIN_ALPHA: u8 = 125; + const MAX_COLOR: u8 = 250; + + let interesting = (a >= MIN_ALPHA) && !(r > MAX_COLOR && g > MAX_COLOR && b > MAX_COLOR); + + !interesting +} + +impl fmt::Display for Palette { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + use itertools::Itertools; + let color_list = self.palette + .iter() + .map(|rgb| format!("#{:02X}{:02X}{:02X}", rgb[0], rgb[1], rgb[2])) + .join(", "); + + write!(f, "Color Palette {{ {} }}", color_list) + } +} \ No newline at end of file diff --git a/src/vibrancy/settings.rs b/src/vibrancy/settings.rs new file mode 100644 index 0000000..fddb0cb --- /dev/null +++ b/src/vibrancy/settings.rs @@ -0,0 +1,21 @@ +#![allow(dead_code)] + +pub const TARGET_DARK_LUMA: f64 = 0.26; +pub const MAX_DARK_LUMA: f64 = 0.45; + +pub const MIN_LIGHT_LUMA: f64 = 0.55; +pub const TARGET_LIGHT_LUMA: f64 = 0.74; + +pub const MIN_NORMAL_LUMA: f64 = 0.3; +pub const TARGET_NORMAL_LUMA: f64 = 0.5; +pub const MAX_NORMAL_LUMA: f64 = 0.7; + +pub const TARGET_MUTED_SATURATION: f64 = 0.3; +pub const MAX_MUTED_SATURATION: f64 = 0.4; + +pub const TARGET_VIBRANT_SATURATION: f64 = 1.0; +pub const MIN_VIBRANT_SATURATION: f64 = 0.35; + +pub const WEIGHT_SATURATION: f64 = 3.0; +pub const WEIGHT_LUMA: f64 = 6.0; +pub const WEIGHT_POPULATION: f64 = 1.0; \ No newline at end of file diff --git a/src/vibrancy/vibrant.rs b/src/vibrancy/vibrant.rs new file mode 100644 index 0000000..34344b5 --- /dev/null +++ b/src/vibrancy/vibrant.rs @@ -0,0 +1,238 @@ +use std::fmt; +use std::collections::BTreeMap; + +use image::{GenericImage, Pixel, Rgb}; + +use hsl::HSL; +// tf is rust on????? +use crate::vibrancy::settings; +// I have no idea why it's imported like this +use super::palette::Palette; + +/// Vibrancy +/// +/// 6 vibrant colors: primary, dark, light, dark muted and light muted. +#[derive(Debug, Hash, PartialEq, Eq, Default)] +pub struct Vibrancy { + pub primary: Option>, + pub dark: Option>, + pub light: Option>, + pub muted: Option>, + pub dark_muted: Option>, + pub light_muted: Option>, +} + +impl Vibrancy { + /// Create new vibrancy map from an image + pub fn new(image: &G) -> Vibrancy + where P: Sized + Pixel, + G: Sized + GenericImage + { + generate_varation_colors(&Palette::new(image, 256, 10)) + } + + fn color_already_set(&self, color: &Rgb) -> bool { + let color = Some(*color); + self.primary == color || self.dark == color || self.light == color || + self.muted == color || self.dark_muted == color || self.light_muted == color + } + + fn find_color_variation(&self, + palette: &[Rgb], + pixel_counts: &BTreeMap, + luma: &MTM, + saturation: &MTM) + -> Option> { + let mut max = None; + let mut max_value = 0_f64; + + let complete_population = pixel_counts.values().fold(0, |acc, c| acc + c); + + for (index, swatch) in palette.iter().enumerate() { + let HSL {h: _, s, l} = HSL::from_rgb(swatch.channels()); + + if s >= saturation.min && s <= saturation.max && l >= luma.min && l <= luma.max && + !self.color_already_set(swatch) { + let population = *pixel_counts.get(&index).unwrap_or(&0) as f64; + if population == 0_f64 { + continue; + } + let value = create_comparison_value(s, + saturation.target, + l, + luma.target, + population, + complete_population as f64); + if max.is_none() || value > max_value { + max = Some(swatch.clone()); + max_value = value; + } + } + } + + max + } + + // fn fill_empty_swatches(self) { + // if self.primary.is_none() { + // // If we do not have a vibrant color... + // if let Some(dark) = self.dark { + // // ...but we do have a dark vibrant, generate the value by modifying the luma + // let hsl = HSL::from_pixel(&dark).clone() + // hsl.l = settings::TARGET_NORMAL_LUMA; + // } + // } + // } +} + +impl fmt::Display for Vibrancy { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + (write!(f, "Vibrant Colors {{\n"))?; + + macro_rules! display_color { + ($formatter:expr, $name:expr, $color:expr) => { + { + (write!($formatter, "\t"))?; + (write!($formatter, $name))?; + if let Some(c) = $color { + let rgb = c.channels(); + (write!($formatter, + " Color: #{:02X}{:02X}{:02X}\n", + rgb[0], rgb[1], rgb[2] + ))?; + } else { + (write!($formatter, " Color: None\n"))?; + } + } + }; + } + + display_color!(f, "Primary Vibrant", self.primary); + display_color!(f, "Dark Vibrant", self.dark); + display_color!(f, "Light Vibrant", self.light); + display_color!(f, "Muted", self.muted); + display_color!(f, "Dark Muted", self.dark_muted); + display_color!(f, "Light Muted", self.light_muted); + + write!(f, "}}") + } +} + +fn generate_varation_colors(p: &Palette) -> Vibrancy { + let mut vibrancy = Vibrancy::default(); + vibrancy.primary = + vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: settings::MIN_NORMAL_LUMA, + target: settings::TARGET_NORMAL_LUMA, + max: settings::MAX_NORMAL_LUMA, + }, + &MTM { + min: settings::MIN_VIBRANT_SATURATION, + target: settings::TARGET_VIBRANT_SATURATION, + max: 1_f64, + }); + + vibrancy.light = vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: settings::MIN_LIGHT_LUMA, + target: settings::TARGET_LIGHT_LUMA, + max: 1_f64, + }, + &MTM { + min: settings::MIN_VIBRANT_SATURATION, + target: settings::TARGET_VIBRANT_SATURATION, + max: 1_f64, + }); + + vibrancy.dark = vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: 0_f64, + target: settings::TARGET_DARK_LUMA, + max: settings::MAX_DARK_LUMA, + }, + &MTM { + min: settings::MIN_VIBRANT_SATURATION, + target: settings::TARGET_VIBRANT_SATURATION, + max: 1_f64, + }); + + vibrancy.muted = vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: settings::MIN_NORMAL_LUMA, + target: settings::TARGET_NORMAL_LUMA, + max: settings::MAX_NORMAL_LUMA, + }, + &MTM { + min: 0_f64, + target: settings::TARGET_MUTED_SATURATION, + max: settings::MAX_MUTED_SATURATION, + }); + + vibrancy.light_muted = vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: settings::MIN_LIGHT_LUMA, + target: settings::TARGET_LIGHT_LUMA, + max: 1_f64, + }, + &MTM { + min: 0_f64, + target: settings::TARGET_MUTED_SATURATION, + max: settings::MAX_MUTED_SATURATION, + }); + + vibrancy.dark_muted = vibrancy.find_color_variation(&p.palette, + &p.pixel_counts, + &MTM { + min: 0_f64, + target: settings::TARGET_DARK_LUMA, + max: settings::MAX_DARK_LUMA, + }, + &MTM { + min: 0_f64, + target: settings::TARGET_MUTED_SATURATION, + max: settings::MAX_MUTED_SATURATION, + }); + + vibrancy +} + +fn invert_diff(val: f64, target_val: f64) -> f64 { + 1_f64 - (val - target_val).abs() +} + +fn weighted_mean(vals: &[(f64, f64)]) -> f64 { + let (sum, sum_weight) = vals.iter().fold((0_f64, 0_f64), + |(sum, sum_weight), &(val, weight)| { + (sum + val * weight, sum_weight + weight) + }); + + sum / sum_weight +} + +fn create_comparison_value(sat: f64, + target_sat: f64, + luma: f64, + target_uma: f64, + population: f64, + max_population: f64) + -> f64 { + weighted_mean(&[(invert_diff(sat, target_sat), + settings::WEIGHT_SATURATION), + (invert_diff(luma, target_uma), settings::WEIGHT_LUMA), + (population / max_population, + settings::WEIGHT_POPULATION)]) +} + +/// Minimum, Maximum, Target +#[derive(Debug, Hash)] +struct MTM { + min: T, + target: T, + max: T, +} \ No newline at end of file