From b3fca8d30c52b9c18b277b2acb07d8bf7949d0d3 Mon Sep 17 00:00:00 2001
From: Adrian Hill <adrian.hill@mailbox.org>
Date: Wed, 6 Mar 2024 22:23:00 +0100
Subject: [PATCH] Allow alpha maps in `heatmap_overlay` (#11)

---
 CHANGELOG.md                             |  5 ++++-
 Project.toml                             |  2 +-
 src/overlay.jl                           | 10 ++++++----
 test/references/overlay_alpha/matrix.txt |  3 +++
 test/references/overlay_alpha/scalar.txt |  3 +++
 test/test_heatmap.jl                     | 21 +++++++++++++++++++++
 6 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 test/references/overlay_alpha/matrix.txt
 create mode 100644 test/references/overlay_alpha/scalar.txt

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a62e783..4da5e22 100644
--- 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
index 7e9ca6b..3209228 100644
--- 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
index 7423d0c..1ebb99f 100644
--- 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
new file mode 100644
index 0000000..6adbb0b
--- /dev/null
+++ b/test/references/overlay_alpha/matrix.txt
@@ -0,0 +1,3 @@
+▀▀▀▀▀▀
+▀▀▀▀▀▀
+▀▀▀▀▀▀
\ No newline at end of file
diff --git a/test/references/overlay_alpha/scalar.txt b/test/references/overlay_alpha/scalar.txt
new file mode 100644
index 0000000..2d1a595
--- /dev/null
+++ b/test/references/overlay_alpha/scalar.txt
@@ -0,0 +1,3 @@
+▀▀▀▀▀▀
+▀▀▀▀▀▀
+▀▀▀▀▀▀
\ No newline at end of file
diff --git a/test/test_heatmap.jl b/test/test_heatmap.jl
index 7af81b9..e4dcf3a 100644
--- 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