Access color schemes through symbols (#3)

Julia-XAI · Oct 18, 2023 · 3df1717 · 3df1717
1 parent c7ec885
commit 3df1717
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/src/VisionHeatmaps.jl b/src/VisionHeatmaps.jl
@@ -1,6 +1,6 @@
 module VisionHeatmaps
 
-using ColorSchemes: ColorScheme, get, seismic
+using ColorSchemes: ColorScheme, colorschemes, get, seismic
 using ImageCore
 
 include("heatmap.jl")

diff --git a/src/heatmap.jl b/src/heatmap.jl
@@ -9,7 +9,7 @@ Visualize 4D arrays as heatmaps, assuming the WHCN convention for input array di
 (width, height, color channels, batch dimension).
 
 ## Keyword arguments
-- `colorscheme::ColorScheme`: Color scheme from ColorSchemes.jl.
+- `colorscheme::Union{ColorScheme,Symbol}`: Color scheme from ColorSchemes.jl.
   Defaults to `seismic`.
 - `reduce::Symbol`: Selects how color channels are reduced to a single number to apply a color scheme.
   The following methods can be selected, which are then applied over the color channels
@@ -32,14 +32,15 @@ Visualize 4D arrays as heatmaps, assuming the WHCN convention for input array di
 """
 function heatmap(
     val::AbstractArray{T,N};
-    colorscheme::ColorScheme=DEFAULT_COLORSCHEME,
+    colorscheme::Union{ColorScheme,Symbol}=DEFAULT_COLORSCHEME,
     reduce::Symbol=DEFAULT_REDUCE,
     rangescale::Symbol=DEFAULT_RANGESCALE,
     permute::Bool=true,
     unpack_singleton::Bool=true,
     process_batch::Bool=false,
 ) where {T,N}
     N != 4 && throw(InputDimensionError)
+    colorscheme = get_colorscheme(colorscheme)
     if unpack_singleton && size(val, 4) == 1
         return single_heatmap(val[:, :, :, 1], colorscheme, reduce, rangescale, permute)
     end
@@ -58,6 +59,9 @@ const InputDimensionError = ArgumentError(
     Please reshape your input to match this format if your model doesn't adhere to this convention.",
 )
 
+get_colorscheme(c::ColorScheme) = c
+get_colorscheme(s::Symbol)::ColorScheme = colorschemes[s]
+
 # Lower level function, mapped along batch dimension
 function single_heatmap(
     val, colorscheme::ColorScheme, reduce::Symbol, rangescale::Symbol, permute::Bool

diff --git a/test/test_heatmap.jl b/test/test_heatmap.jl
@@ -30,6 +30,10 @@ end
 @testset "ColorSchemes" begin
     h = heatmap(A; colorscheme=ColorSchemes.inferno)
     @test_reference "references/inferno.txt" h
+
+    # Test colorscheme symbols
+    h = heatmap(A; colorscheme=:inferno)
+    @test_reference "references/inferno.txt" h
 end
 
 @testset "Error handling" begin