From 1cf012c6e47f53fb8443e9073731441911f33c1e Mon Sep 17 00:00:00 2001 From: Simon Niedermayr Date: Thu, 16 May 2024 10:46:42 +0200 Subject: [PATCH] better cmaps --- colormaps.py | 92 ++++++++++++++----------- requirements.txt | 4 +- src/bin/render.rs | 4 +- src/bin/viewer.rs | 4 +- src/cmap.rs | 171 ++++++++++++++++++++++++++++++++++------------ src/offline.rs | 4 +- src/py.rs | 2 +- src/ui.rs | 106 +++++++++++++++++++++------- src/web.rs | 6 +- 9 files changed, 276 insertions(+), 117 deletions(-) diff --git a/colormaps.py b/colormaps.py index 4c7b1f6..7855b93 100644 --- a/colormaps.py +++ b/colormaps.py @@ -2,46 +2,62 @@ import json from matplotlib import pyplot as plt from os import makedirs -from matplotlib.colors import LinearSegmentedColormap, ListedColormap +from matplotlib.colors import LinearSegmentedColormap, ListedColormap, Colormap import numpy as np -makedirs("colormaps", exist_ok=True) -for name in plt.colormaps(): - cmap = plt.get_cmap(name) - if isinstance(cmap, ListedColormap): - data = cmap(np.linspace(0, 1, 256)).astype(np.float32) - np.save("colormaps/{}.npy".format(name), data) - elif isinstance(cmap, LinearSegmentedColormap): - segments = cmap._segmentdata - - if any( - not isinstance(segments[c], np.ndarray) - and not isinstance(segments[c], list) - and not isinstance(segments[c], tuple) - for c in ["red", "green", "blue"] - ): + +def export_colormaps(cm_module, folder: str): + makedirs(folder, exist_ok=True) + for name, cmap in vars(cm_module).items(): + if not isinstance(cmap, Colormap): + continue + if name.endswith("_r"): + continue + if name.startswith("cmr."): + name = name[4:] + if isinstance(cmap, ListedColormap): data = cmap(np.linspace(0, 1, 256)).astype(np.float32) - np.save("colormaps/{}.npy".format(name), data) - else: - channels = { - c: ( - segments[c].tolist() - if isinstance(segments[c], np.ndarray) - else segments[c] - ) + np.save("{}/{}.npy".format(folder, name), data) + elif isinstance(cmap, LinearSegmentedColormap): + segments = cmap._segmentdata + + if any( + not isinstance(segments[c], np.ndarray) + and not isinstance(segments[c], list) + and not isinstance(segments[c], tuple) for c in ["red", "green", "blue"] - } - if "alpha" in channels: - channels["alpha"] = segments["alpha"].tolist() - with open("colormaps/{}.json".format(name), "w") as f: - json.dump( - channels, - f, - indent=2, - ) - else: - raise TypeError( - "cmap must be either a ListedColormap or a LinearSegmentedColormap" - ) + ): + data = cmap(np.linspace(0, 1, 256)).astype(np.float32) + np.save("{}/{}.npy".format(folder, name), data) + else: + channels = { + c: ( + segments[c].tolist() + if isinstance(segments[c], np.ndarray) + else segments[c] + ) + for c in ["red", "green", "blue"] + } + if "alpha" in channels: + channels["alpha"] = segments["alpha"].tolist() + with open("{}/{}.json".format(folder, name), "w") as f: + json.dump( + channels, + f, + indent=2, + ) + else: + raise TypeError( + "cmap must be either a ListedColormap or a LinearSegmentedColormap" + ) -# %% + +export_colormaps(plt.cm, "colormaps/matplotlib") + +import seaborn as sns + +export_colormaps(sns.cm, "colormaps/seaborn") + +import cmasher + +export_colormaps(cmasher.cm, "colormaps/cmasher") diff --git a/requirements.txt b/requirements.txt index 8248f97..b510a31 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ matplotlib -numpy \ No newline at end of file +numpy +seaborn +cmasher \ No newline at end of file diff --git a/src/bin/render.rs b/src/bin/render.rs index 113b404..3367848 100644 --- a/src/bin/render.rs +++ b/src/bin/render.rs @@ -2,7 +2,7 @@ use cgmath::Vector2; use clap::Parser; use std::{fs::File, path::PathBuf}; -use v4dv::cmap::ColorMapType; +use v4dv::cmap::GenericColorMap; use v4dv::volume::Volume; #[derive(Debug, Parser)] @@ -54,7 +54,7 @@ async fn main() -> anyhow::Result<()> { let volume_file = File::open(&opt.input)?; let cmap_file = File::open(opt.colormap)?; let volumes = Volume::load_numpy(volume_file, opt.channel_first)?; - let cmap = ColorMapType::read(cmap_file)?; + let cmap = GenericColorMap::read(cmap_file)?; let resolution = Vector2::new(opt.width, opt.height); diff --git a/src/bin/viewer.rs b/src/bin/viewer.rs index 6d63cb9..ed5490a 100644 --- a/src/bin/viewer.rs +++ b/src/bin/viewer.rs @@ -34,11 +34,11 @@ async fn main() -> anyhow::Result<()> { .expect("Failed to load volume"); #[cfg(feature = "colormaps")] - let cmap = cmap::COLORMAPS.get("summer").unwrap().clone(); + let cmap = cmap::COLORMAPS["matplotlib"]["viridis"].clone(); #[cfg(not(feature = "colormaps"))] let cmap = { let reader = File::open(&opt.colormap)?; - cmap::ColorMapType::read(reader)? + cmap::GenericColorMap::read(reader)? }; open_window( diff --git a/src/cmap.rs b/src/cmap.rs index fa7f4dc..9607e85 100644 --- a/src/cmap.rs +++ b/src/cmap.rs @@ -8,6 +8,8 @@ use std::{collections::HashMap, io::Cursor}; use anyhow::Ok; use cgmath::Vector4; +#[cfg(feature = "colormaps")] +use include_dir::Dir; use npyz::WriterBuilder; #[cfg(feature = "python")] use numpy::ndarray::{ArrayViewD, Axis}; @@ -19,20 +21,37 @@ use once_cell::sync::Lazy; #[cfg(feature = "colormaps")] use include_dir::include_dir; -// list of predefined colormaps #[cfg(feature = "colormaps")] -pub static COLORMAPS: Lazy> = Lazy::new(|| { - let cmaps: HashMap = include_dir!("colormaps") +static COLORMAPS_MATPLOTLIB: include_dir::Dir = include_dir!("colormaps/matplotlib"); +#[cfg(feature = "colormaps")] +static COLORMAPS_SEABORN: include_dir::Dir = include_dir!("colormaps/seaborn"); +#[cfg(feature = "colormaps")] +static COLORMAPS_CMASHER: include_dir::Dir = include_dir!("colormaps/cmasher"); + +#[cfg(feature = "colormaps")] +fn load_cmaps(dir: &Dir) -> HashMap { + let cmaps: HashMap = dir .files() .filter_map(|f| { let file_name = f.path(); let reader = Cursor::new(f.contents()); let name = file_name.file_stem().unwrap().to_str().unwrap().to_string(); - let cmap = ColorMapType::read(reader).unwrap(); + let cmap = GenericColorMap::read(reader).unwrap(); return Some((name, cmap)); }) .collect(); cmaps +} + +// list of predefined colormaps +#[cfg(feature = "colormaps")] +pub static COLORMAPS: Lazy>> = Lazy::new(|| { + let mut cmaps = HashMap::new(); + cmaps.insert("matplotlib".to_string(), load_cmaps(&COLORMAPS_MATPLOTLIB)); + cmaps.insert("seaborn".to_string(), load_cmaps(&COLORMAPS_SEABORN)); + cmaps.insert("cmasher".to_string(), load_cmaps(&COLORMAPS_CMASHER)); + + cmaps }); #[derive(Debug, Clone)] @@ -116,26 +135,19 @@ impl ListedColorMap { } } -// i dont no why we need this but it works... -impl ColorMap for &dyn ColorMap { - fn sample(&self, x: f32) -> Vector4 { - (*self).sample(x) - } -} - -impl ColorMap for ListedColorMap { - fn sample(&self, x: f32) -> Vector4 { - let n = self.0.len() as f32; - let i = (x * n).min(n - 1.0).max(0.0) as usize; - self.0[i] // TODO linear interpolation - } -} impl<'a> ColorMap for &'a ListedColorMap { + type Item = ListedColorMap; fn sample(&self, x: f32) -> Vector4 { let n = self.0.len() as f32; let i = (x * n).min(n - 1.0).max(0.0) as usize; self.0[i] // TODO linear interpolation } + + fn reverse(&self) -> ListedColorMap { + let mut cmap = self.0.clone(); + cmap.reverse(); + ListedColorMap::new(cmap) + } } pub struct ColorMapGPU { @@ -144,6 +156,7 @@ pub struct ColorMapGPU { } pub trait ColorMap { + type Item; fn sample(&self, x: f32) -> Vector4; fn rasterize(&self, n: usize) -> Vec> { @@ -151,6 +164,8 @@ pub trait ColorMap { .map(|i| self.sample(i as f32 / (n - 1) as f32)) .collect() } + + fn reverse(&self) -> Self::Item; } impl ColorMapGPU { @@ -279,12 +294,8 @@ pub struct LinearSegmentedColorMap { pub g: Vec<(f32, f32, f32)>, #[serde(alias = "blue")] pub b: Vec<(f32, f32, f32)>, - #[serde(alias = "alpha", default = "default_alpha")] - pub a: Vec<(f32, f32, f32)>, -} - -fn default_alpha() -> Vec<(f32, f32, f32)> { - vec![(0., 1., 1.), (1., 1., 1.)] + #[serde(alias = "alpha")] + pub a: Option>, } impl LinearSegmentedColorMap { @@ -292,7 +303,7 @@ impl LinearSegmentedColorMap { r: Vec<(f32, f32, f32)>, g: Vec<(f32, f32, f32)>, b: Vec<(f32, f32, f32)>, - a: Vec<(f32, f32, f32)>, + a: Option>, ) -> anyhow::Result { if !Self::check_values(&r) { return Err(anyhow::anyhow!( @@ -302,21 +313,23 @@ impl LinearSegmentedColorMap { if !Self::check_values(&g) { return Err(anyhow::anyhow!( - "x values for red are not in (0,1) or ascending" + "x values for green are not in (0,1) or ascending" )); }; if !Self::check_values(&b) { return Err(anyhow::anyhow!( - "x values for red are not in (0,1) or ascending" + "x values for blue are not in (0,1) or ascending" )); }; - if !Self::check_values(&a) { - return Err(anyhow::anyhow!( - "x values for red are not in (0,1) or ascending" - )); - }; + if let Some(a) = &a { + if !Self::check_values(&a) { + return Err(anyhow::anyhow!( + "x values for alpha are not in (0,1) or ascending" + )); + }; + } Ok(Self { r, g, b, a }) } @@ -353,29 +366,75 @@ impl LinearSegmentedColorMap { merge_neighbours(&mut b); merge_neighbours(&mut a); - Self { r, g, b, a } + Self { + r, + g, + b, + a: Some(a), + } } } impl ColorMap for &LinearSegmentedColorMap { + type Item = LinearSegmentedColorMap; fn sample(&self, x: f32) -> Vector4 { + let a = self + .a + .as_ref() + .map(|a| sample_channel(x, &a)) + .unwrap_or(1.0); Vector4::new( (sample_channel(x, &self.r) * 255.) as u8, (sample_channel(x, &self.g) * 255.) as u8, (sample_channel(x, &self.b) * 255.) as u8, - (sample_channel(x, &self.a) * 255.) as u8, + (a * 255.) as u8, ) } + fn reverse(&self) -> Self::Item { + let mut r: Vec<_> = self + .r + .iter() + .map(|(x, y1, y2)| (1.0 - x, *y1, *y2)) + .collect(); + let mut g: Vec<_> = self + .g + .iter() + .map(|(x, y1, y2)| (1.0 - x, *y1, *y2)) + .collect(); + let mut b: Vec<_> = self + .b + .iter() + .map(|(x, y1, y2)| (1.0 - x, *y1, *y2)) + .collect(); + let mut a: Option> = self + .a + .clone() + .map(|a| a.iter().map(|(x, y1, y2)| (1.0 - x, *y1, *y2)).collect()); + r.reverse(); + g.reverse(); + b.reverse(); + if let Some(a) = &mut a { + a.reverse(); + } + LinearSegmentedColorMap { r, g, b, a } + } } impl Hash for LinearSegmentedColorMap { fn hash(&self, state: &mut H) { - for c in [&self.r, &self.g, &self.b, &self.a].iter() { + for c in [&self.r, &self.g, &self.b].iter() { c.iter().for_each(|(a, b, c)| { state.write_u32(a.to_bits()); state.write_u32(b.to_bits()); state.write_u32(c.to_bits()) }); } + if let Some(a) = &self.a { + a.iter().for_each(|(a, b, c)| { + state.write_u32(a.to_bits()); + state.write_u32(b.to_bits()); + state.write_u32(c.to_bits()) + }); + } } } @@ -434,22 +493,22 @@ fn merge_neighbours(values: &mut Vec<(f32, f32, f32)>) { pub const COLORMAP_RESOLUTION: u32 = 256; #[derive(Debug, Clone, Hash)] -pub enum ColorMapType { +pub enum GenericColorMap { Listed(ListedColorMap), LinearSegmented(LinearSegmentedColorMap), } -impl ColorMapType { +impl GenericColorMap { pub fn read(mut reader: R) -> anyhow::Result { let mut start = [0; 6]; reader.read_exact(&mut start)?; reader.seek(SeekFrom::Start(0))?; if start.eq(b"\x93NUMPY") { // numpy file - Ok(ColorMapType::Listed(ListedColorMap::from_npy(reader)?)) + Ok(GenericColorMap::Listed(ListedColorMap::from_npy(reader)?)) } else { // json file - Ok(ColorMapType::LinearSegmented( + Ok(GenericColorMap::LinearSegmented( LinearSegmentedColorMap::from_json(reader)?, )) } @@ -457,17 +516,41 @@ impl ColorMapType { pub fn into_linear_segmented(&self, n: u32) -> LinearSegmentedColorMap { match self { - crate::cmap::ColorMapType::Listed(c) => LinearSegmentedColorMap::from_color_map(c, n), - crate::cmap::ColorMapType::LinearSegmented(c) => c.clone(), + crate::cmap::GenericColorMap::Listed(c) => { + LinearSegmentedColorMap::from_color_map(c, n) + } + crate::cmap::GenericColorMap::LinearSegmented(c) => c.clone(), + } + } + + /// if all alpha values are 1.0, the alpha channel is considered boring + #[allow(unused)] + pub(crate) fn has_boring_alpha_channel(&self) -> bool { + match self { + GenericColorMap::Listed(c) => c.0.iter().all(|v| v.w == 255), + GenericColorMap::LinearSegmented(c) => { + c.a.as_ref() + .map(|a| a.iter().all(|(_, _, a)| *a == 1.0)) + .unwrap_or(true) + } } } } -impl<'a> ColorMap for &'a ColorMapType { +impl<'a> ColorMap for &'a GenericColorMap { fn sample(&self, x: f32) -> Vector4 { match self { - ColorMapType::Listed(c) => c.sample(x), - ColorMapType::LinearSegmented(c) => c.sample(x), + GenericColorMap::Listed(c) => c.sample(x), + GenericColorMap::LinearSegmented(c) => c.sample(x), + } + } + + type Item = GenericColorMap; + + fn reverse(&self) -> Self::Item { + match self { + GenericColorMap::Listed(c) => GenericColorMap::Listed(c.reverse()), + GenericColorMap::LinearSegmented(c) => GenericColorMap::LinearSegmented(c.reverse()), } } } diff --git a/src/offline.rs b/src/offline.rs index 83bf3ff..c2e6bc1 100644 --- a/src/offline.rs +++ b/src/offline.rs @@ -3,7 +3,7 @@ use image::{ImageBuffer, Rgba}; use crate::{ camera::{GenericCamera, PerspectiveProjection, Projection}, - cmap::{ColorMapGPU, ColorMapType, COLORMAP_RESOLUTION}, + cmap::{ColorMapGPU, GenericColorMap, COLORMAP_RESOLUTION}, renderer::{RenderSettings, VolumeRenderer}, volume::{Volume, VolumeGPU}, WGPUContext, @@ -65,7 +65,7 @@ async fn render_view( pub async fn render_volume( volumes: Vec, - cmap: ColorMapType, + cmap: GenericColorMap, resolution: Vector2, time: f32, bg: wgpu::Color, diff --git a/src/py.rs b/src/py.rs index 9101cdd..a43ce87 100644 --- a/src/py.rs +++ b/src/py.rs @@ -28,7 +28,7 @@ fn v4dv<'py>(m: &Bound<'py, PyModule>) -> PyResult<()> { let cmap = ListedColorMap::from_array(cmap.as_array()); let img = pollster::block_on(render_volume( vec![volume], - cmap::ColorMapType::Listed(cmap), + cmap::GenericColorMap::Listed(cmap), Vector2::new(width, height), time, wgpu::Color { diff --git a/src/ui.rs b/src/ui.rs index 544c000..1ac7c65 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -147,38 +147,74 @@ pub(crate) fn ui(state: &mut WindowContext) { if state.cmap_select_visible { ui.horizontal(|ui| { let cmaps = &COLORMAPS; - let mut selected_cmap: String = ui.ctx().data_mut(|d| { - d.get_persisted_mut_or("selected_cmap".into(), "viridis".to_string()) + let mut selected_cmap: (String, String) = ui.ctx().data_mut(|d| { + d.get_persisted_mut_or( + "selected_cmap".into(), + ("matplotlib".to_string(), "viridis".to_string()), + ) + .clone() + }); + let mut search_term: String = ui.ctx().data_mut(|d| { + d.get_temp_mut_or("cmap_search".into(), "".to_string()) .clone() }); ui.label("Colormap"); let old_selected_cmap = selected_cmap.clone(); egui::ComboBox::new("cmap_select", "") - .selected_text(selected_cmap.clone()) + .selected_text(selected_cmap.1.clone()) .show_ui(ui, |ui| { - let mut keys: Vec<_> = cmaps.iter().collect(); - keys.sort_by_key(|e| e.0); - for (name, cmap) in keys { - let texture = load_or_create(ui, cmap, COLORMAP_RESOLUTION); - ui.horizontal(|ui| { - ui.image(egui::ImageSource::Texture( - egui::load::SizedTexture { - id: texture, - size: vec2(50., 10.), - }, - )); - ui.selectable_value(&mut selected_cmap, name.clone(), name); - }); + ui.add( + egui::text_edit::TextEdit::singleline(&mut search_term) + .hint_text("Search..."), + ); + for (group, cmaps) in cmaps.iter() { + ui.label(group); + let mut sorted_cmaps: Vec<_> = cmaps.iter().collect(); + sorted_cmaps.sort_by_key(|e| e.0); + for (name, cmap) in sorted_cmaps { + if name.contains(&search_term) { + let texture = + load_or_create(ui, cmap, COLORMAP_RESOLUTION); + ui.horizontal(|ui| { + ui.image(egui::ImageSource::Texture( + egui::load::SizedTexture { + id: texture, + size: vec2(50., 10.), + }, + )); + ui.selectable_value( + &mut selected_cmap, + (group.clone(), name.clone()), + name, + ); + }); + } + } + ui.separator(); } }); if old_selected_cmap != selected_cmap { - state.cmap = - cmaps[&selected_cmap].into_linear_segmented(COLORMAP_RESOLUTION); - + let old_alpha = state.cmap.a.clone(); + state.cmap = cmaps[&selected_cmap.0][&selected_cmap.1] + .into_linear_segmented(COLORMAP_RESOLUTION); + if state.cmap.a.is_none() + || cmaps[&selected_cmap.0][&selected_cmap.1] + .has_boring_alpha_channel() + { + state.cmap.a = old_alpha; + } ui.ctx().data_mut(|d| { - d.insert_persisted("selected_cmap".into(), selected_cmap) + d.insert_persisted("selected_cmap".into(), selected_cmap); }); } + ui.ctx() + .data_mut(|d| d.insert_temp("cmap_search".into(), search_term)); + if state.cmap.a.is_none() { + state.cmap.a = Some(vec![(0.0, 1.0, 1.0), (1.0, 1.0, 1.0)]); + } + if ui.button("↔").clicked() { + state.cmap = (&state.cmap).reverse(); + } }); } let vmin = state @@ -192,11 +228,33 @@ pub(crate) fn ui(state: &mut WindowContext) { show_cmap(ui, egui::Id::new("cmap preview"), &state.cmap, vmin, vmax); ui.label("Alpha Channel"); - if ui.button("felix hack").clicked() { - state.cmap.a = vec![(0.0, 1.0, 1.0), (0.5, 0., 0.), (1.0, 1.0, 1.0)]; - } + ui.end_row(); + ui.horizontal(|ui| { + ui.label("Presets:"); + if ui.button("felix hack").clicked() { + state.cmap.a = Some(vec![(0.0, 1.0, 1.0), (0.5, 0., 0.), (1.0, 1.0, 1.0)]); + } + if ui.button("simon hack").clicked() { + state.cmap.a = Some(vec![(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]); + } + if ui.button("double felix hack").clicked() { + state.cmap.a = Some(vec![ + (0.0, 0.0, 0.0), + (0.25, 1.0, 1.0), + (0.5, 0.0, 0.0), + (0.75, 1.0, 1.0), + (1.0, 0.0, 0.0), + ]); + } + if ui.button("flat").clicked() { + state.cmap.a = Some(vec![(0.0, 1.0, 1.0), (1.0, 1.0, 1.0)]); + } + }); + ui.separator(); - tf_ui(ui, &mut state.cmap.a); + if let Some(a) = &mut state.cmap.a { + tf_ui(ui, a); + } ui.end_row(); #[cfg(not(target_arch = "wasm32"))] ui.horizontal(|ui| { diff --git a/src/web.rs b/src/web.rs index 5e6bc32..707bf32 100644 --- a/src/web.rs +++ b/src/web.rs @@ -3,7 +3,7 @@ use wasm_bindgen::JsCast; use winit::platform::web::WindowBuilderExtWebSys; use winit::window::WindowBuilder; -use crate::cmap::{ColorMapType, COLORMAP_RESOLUTION}; +use crate::cmap::{GenericColorMap, COLORMAP_RESOLUTION}; use crate::volume::Volume; use crate::{open_window, RenderConfig}; @@ -87,7 +87,7 @@ pub async fn viewer_inline( let reader_colormap = Cursor::new(colormap); - let cmap = ColorMapType::read(reader_colormap) + let cmap = GenericColorMap::read(reader_colormap) .unwrap() .into_linear_segmented(COLORMAP_RESOLUTION); @@ -140,7 +140,7 @@ pub async fn viewer_wasm(canvas_id: String) { std::panic::set_hook(Box::new(console_error_panic_hook::hook)); console_log::init().expect("could not initialize logger"); - let cmap = cmap::COLORMAPS.get("viridis").unwrap().clone(); + let cmap = cmap::COLORMAPS["seaborn"]["icefire"].clone(); let (canvas, spinner): (HtmlCanvasElement, HtmlElement) = web_sys::window() .and_then(|win| win.document())