Allow alpha maps in heatmap_overlay (#11)

Julia-XAI · Mar 6, 2024 · b3fca8d · b3fca8d
1 parent 2cacfed
commit b3fca8d
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,9 +1,11 @@
 # VisionHeatmaps.jl
+## Version `v1.5.0`
+* ![Feature][badge-feature] Allow alpha maps in `heatmap_overlay` ([#11][pr-11])
+
 ## Version `v1.4.0`
 * ![Feature][badge-feature] Add heatmap overlays on Explanations ([#10][pr-10])
 * ![Maintenance][badge-maintenance] Use Configurations.jl for keyword argument handling ([#9][pr-9])
 
-
 ## Version `v1.3.1`
 * ![Feature][badge-feature] Add XAIBase dependency ([#7][pr-7], [#8][pr-8])
 
@@ -17,6 +19,7 @@
 ## Version `v1.0.0`
 * Initial release
 
+[pr-11]: https://github.com/Julia-XAI/VisionHeatmaps.jl/pull/11
 [pr-10]: https://github.com/Julia-XAI/VisionHeatmaps.jl/pull/10
 [pr-9]: https://github.com/Julia-XAI/VisionHeatmaps.jl/pull/9
 [pr-8]: https://github.com/Julia-XAI/VisionHeatmaps.jl/pull/8

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "VisionHeatmaps"
 uuid = "27106da1-f8bc-4ca8-8c66-9b8289f1e035"
 authors = ["Adrian Hill <gh@adrianhill.de>"]
-version = "1.4.0"
+version = "1.5.0-DEV"
 
 [deps]
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"

diff --git a/src/overlay.jl b/src/overlay.jl
@@ -11,7 +11,9 @@ Assumes 4D input array following the WHCN convention
 (width, height, color channels, batch dimension) and batch size 1.
 
 ## Keyword arguments
-- `alpha::Real`: Opacity of the heatmap overlay. Defaults to `$DEFAULT_OVERLAY_ALPHA`.
+- `alpha`: Opacity of the heatmap overlay. Defaults to `$DEFAULT_OVERLAY_ALPHA`.
+    If an array is passed, it will be broadcast with the image and heatmap.
+    All alpha values are expected to be between 0 and 1.
 - `resize_method`: Method used to resize the heatmap in case of a size mismatch with the image.
     Defaults to `Lanczos(1)` from Interpolations.jl.
 
@@ -32,7 +34,7 @@ end
 function heatmap_overlay(
     val::AbstractArray{T,N},
     im::AbstractImage,
-    alpha::Real,
+    alpha,
     resize_method,
     options::HeatmapOptions,
 ) where {T,N}
@@ -44,7 +46,7 @@ function heatmap_overlay(
             ),
         )
     end
-    if alpha < 0 || alpha > 1
+    if any(<(0), alpha) || any(>(1), alpha)
         throw(ArgumentError("alpha must be in the range [0, 1]"))
     end
 
@@ -54,7 +56,7 @@ function heatmap_overlay(
     if hmsize != imsize
         hm = imresize(hm, imsize; method=resize_method)
     end
-    return im * (1 - alpha) + hm * alpha
+    return @. im * (1 - alpha) + hm * alpha
 end
 
 #=================#

diff --git a/test/references/overlay_alpha/matrix.txt b/test/references/overlay_alpha/matrix.txt
@@ -0,0 +1,3 @@
+[0m[38;5;232;48;5;52m▀[38;5;16;48;5;52m▀[38;5;22;48;5;58m▀[38;5;64;48;5;94m▀[38;5;64;48;5;100m▀[38;5;64;48;5;100m▀[0m
+[0m[38;5;88;48;5;124m▀[38;5;88;48;5;124m▀[38;5;94;48;5;130m▀[38;5;94;48;5;94m▀[38;5;94;48;5;94m▀[38;5;130;48;5;88m▀[0m
+[0m[38;5;124;48;5;160m▀[38;5;124;48;5;124m▀[38;5;124;48;5;124m▀[38;5;124;48;5;124m▀[38;5;88;48;5;88m▀[38;5;88;48;5;88m▀[0m
diff --git a/test/references/overlay_alpha/scalar.txt b/test/references/overlay_alpha/scalar.txt
@@ -0,0 +1,3 @@
+[0m[38;5;16;48;5;52m▀[38;5;16;48;5;52m▀[38;5;22;48;5;58m▀[38;5;28;48;5;64m▀[38;5;34;48;5;70m▀[38;5;40;48;5;76m▀[0m
+[0m[38;5;88;48;5;124m▀[38;5;88;48;5;124m▀[38;5;94;48;5;130m▀[38;5;100;48;5;100m▀[38;5;106;48;5;106m▀[38;5;112;48;5;112m▀[0m
+[0m[38;5;124;48;5;160m▀[38;5;124;48;5;160m▀[38;5;130;48;5;166m▀[38;5;136;48;5;172m▀[38;5;142;48;5;178m▀[38;5;148;48;5;184m▀[0m
diff --git a/test/test_heatmap.jl b/test/test_heatmap.jl
@@ -70,6 +70,27 @@ end
     end
 end
 
+@testset "Overlay alpha masks" begin
+    alpha = 0.2
+    ho = heatmap_overlay(A, img2; alpha=alpha)
+    @test size(ho) == size(img2)
+    @test_reference "references/overlay_alpha/scalar.txt" ho
+
+    alpha = [(x + y) / 2 for x in 0:0.2:1, y in 0:0.2:1]
+    ho = heatmap_overlay(A, img2; alpha=alpha)
+    @test size(ho) == size(img2)
+    @test_reference "references/overlay_alpha/matrix.txt" ho
+
+    alpha = fill!(alpha, 0.2)
+    ho = heatmap_overlay(A, img2; alpha=alpha)
+    @test size(ho) == size(img2)
+    @test_reference "references/overlay_alpha/scalar.txt" ho
+
+    alpha[2, 2] = 1.1
+    @test_throws ArgumentError heatmap_overlay(A, img2; alpha=alpha)
+    @test_throws ArgumentError heatmap_overlay(A, img2; alpha=-0.1)
+end
+
 @testset "Batched input" begin
     for reducer in reducers
         for rangescale in rangescales