Update XAIBase to v3.0.0 and document heatmap_overlay (#162)

This is a breaking change since heatmap now requires an explicit import of either VisionHeatmaps.jl or TextHeatmaps.jl:

* Update XAIBase to v3.0.0

* Import VisionHeatmaps in docs

* Document `heatmap_overlay`

Author: adrhill

Project.toml

@@ -1,7 +1,7 @@
 name = "ExplainableAI"
 uuid = "4f1bc3e1-d60d-4ed0-9367-9bdff9846d3b"
 authors = ["Adrian Hill <[email protected]>"]
-version = "0.7.0"
+version = "0.8.0-DEV"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -16,6 +16,6 @@ Distributions = "0.25"
 Random = "1"
 Reexport = "1"
 Statistics = "1"
-XAIBase = "1.2"
+XAIBase = "3"
 Zygote = "0.6"
 julia = "1.6"

docs/Project.toml

@@ -10,3 +10,7 @@ ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
 ImageShow = "4e3cecfd-b093-5904-9786-8bbb286a6a31"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+VisionHeatmaps = "27106da1-f8bc-4ca8-8c66-9b8289f1e035"
+
+[compat]
+VisionHeatmaps = "1.4"

docs/make.jl

@@ -25,7 +25,11 @@ makedocs(;
     modules=[XAIBase, ExplainableAI],
     authors="Adrian Hill",
     sitename="ExplainableAI.jl",
-    format=Documenter.HTML(; prettyurls=get(ENV, "CI", "false") == "true", assets=String[]),
+    format=Documenter.HTML(;
+        prettyurls=get(ENV, "CI", "false") == "true",
+        size_threshold=300_000,
+        assets=String[],
+    ),
     #! format: off
     pages=[
         "Home" => "index.md",

docs/src/api.md

@@ -3,9 +3,14 @@ All methods in ExplainableAI.jl work by calling `analyze` on an input and an ana
 ```@docs
 analyze
 Explanation
-heatmap
 ```
 
+For heatmapping functionality, take a look at either
+[VisionHeatmaps.jl](https://julia-xai.github.io/XAIDocs/VisionHeatmaps/stable/) or
+[TextHeatmaps.jl](https://julia-xai.github.io/XAIDocs/TextHeatmaps/stable/).
+Both provide `heatmap` methods for visualizing explanations, 
+either for images or text, respectively.
+
 # Analyzers
 ```@docs
 Gradient
@@ -27,3 +32,4 @@ InterpolationAugmentation
 # Index
 ```@index
 ```
+

docs/src/literate/augmentations.jl

@@ -7,6 +7,7 @@
 # We build on the basics shown in the [*Getting started*](@ref docs-getting-started) section
 # and start out by loading the same pre-trained LeNet5 model and MNIST input data:
 using ExplainableAI
+using VisionHeatmaps
 using Flux
 
 using BSON # hide

docs/src/literate/example.jl

@@ -68,7 +68,11 @@ expl.val
 
 # ## Heatmapping basics
 # Since the array `expl.val` is not very informative at first sight,
-# we can visualize `Explanation`s by computing a [`heatmap`](@ref):
+# we can visualize `Explanation`s by computing a `heatmap` using either
+# [VisionHeatmaps.jl](https://julia-xai.github.io/XAIDocs/VisionHeatmaps/stable/) or
+# [TextHeatmaps.jl](https://julia-xai.github.io/XAIDocs/TextHeatmaps/stable/).
+using VisionHeatmaps
+
 heatmap(expl)
 
 # If we are only interested in the heatmap, we can combine analysis and heatmapping
@@ -92,7 +96,7 @@ heatmap(expl)
 
 #md # !!! note
 #md #
-#md #     The output neuron can also be specified when calling [`heatmap`](@ref):
+#md #     The output neuron can also be specified when calling `heatmap`:
 #md #     ```julia
 #md #     heatmap(input, analyzer, 5)
 #md #     ```

docs/src/literate/heatmapping.jl

@@ -1,11 +1,15 @@
 # # [Heatmapping](@id docs-heatmapping)
 # Since numerical explanations are not very informative at first sight,
-# we can visualize them by computing a [`heatmap`](@ref).
+# we can visualize them by computing a `heatmap`, using either
+# [VisionHeatmaps.jl](https://julia-xai.github.io/XAIDocs/VisionHeatmaps/stable/) or
+# [TextHeatmaps.jl](https://julia-xai.github.io/XAIDocs/TextHeatmaps/stable/).
+#
 # This page showcases different options and preset for heatmapping,
 # building on the basics shown in the [*Getting started*](@ref docs-getting-started) section.
 #
 # We start out by loading the same pre-trained LeNet5 model and MNIST input data:
 using ExplainableAI
+using VisionHeatmaps
 using Flux
 
 using BSON # hide
@@ -19,10 +23,10 @@ index = 10
 x, y = MNIST(Float32, :test)[10]
 input = reshape(x, 28, 28, 1, :)
 
-convert2image(MNIST, x)
+img = convert2image(MNIST, x)
 
 # ## Automatic heatmap presets
-# The function [`heatmap`](@ref) automatically applies common presets for each method.
+# The function `heatmap` automatically applies common presets for each method.
 #
 # Since [`InputTimesGradient`](@ref) computes attributions,
 # heatmaps are shown in a blue-white-red color scheme.
@@ -35,7 +39,7 @@ heatmap(input, analyzer)
 
 # ## Custom heatmap settings
 # ### Color schemes
-# We can partially or fully override presets by passing keyword arguments to [`heatmap`](@ref).
+# We can partially or fully override presets by passing keyword arguments to `heatmap`.
 # For example, we can use a custom color scheme from ColorSchemes.jl using the keyword argument `colorscheme`:
 using ColorSchemes
 
@@ -89,20 +93,32 @@ heatmap(expl; rangescale=:centered, colorscheme=:inferno)
 #-
 heatmap(expl; rangescale=:extrema, colorscheme=:inferno)
 
-# For the full list of `heatmap` keyword arguments, refer to the [`heatmap`](@ref) documentation.
+# For the full list of `heatmap` keyword arguments, refer to the `heatmap` documentation.
+
+# ## [Heatmap overlays](@id overlay)
+# Heatmaps can be overlaid onto the input image using the `heatmap_overlay` function
+# from VisionHeatmaps.jl.
+# This can be useful for visualizing the relevance of specific regions of the input:
+heatmap_overlay(expl, img)
+
+# The alpha value of the heatmap can be adjusted using the `alpha` keyword argument:
+heatmap_overlay(expl, img; alpha=0.3)
+
+# All previously discussed keyword arguments for `heatmap` can also be used with `heatmap_overlay`:
+heatmap_overlay(expl, img; alpha=0.7, colorscheme=:inferno, rangescale=:extrema)
 
 # ## [Heatmapping batches](@id docs-heatmapping-batches)
 # Heatmapping also works with input batches.
-# Let's demonstrate this by using a batch of 100 images from the MNIST dataset:
-xs, ys = MNIST(Float32, :test)[1:100]
+# Let's demonstrate this by using a batch of 25 images from the MNIST dataset:
+xs, ys = MNIST(Float32, :test)[1:25]
 batch = reshape(xs, 28, 28, 1, :); # reshape to WHCN format
 
-# The [`heatmap`](@ref) function automatically recognizes
+# The `heatmap` function automatically recognizes
 # that the explanation is batched and returns a `Vector` of images:
 heatmaps = heatmap(batch, analyzer)
 
 # Image.jl's `mosaic` function can used to display them in a grid:
-mosaic(heatmaps; nrow=10)
+mosaic(heatmaps; nrow=5)
 
 # When heatmapping batches, the mapping to the color scheme is applied per sample.
 # For example, `rangescale=:extrema` will normalize each heatmap
@@ -113,12 +129,12 @@ mosaic(heatmaps; nrow=10)
 # `heatmap` can be called with the keyword-argument `process_batch=true`:
 expl = analyze(batch, analyzer)
 heatmaps = heatmap(expl; process_batch=true)
-mosaic(heatmaps; nrow=10)
+mosaic(heatmaps; nrow=5)
 
 # This can be useful when comparing heatmaps for fixed output neurons:
 expl = analyze(batch, analyzer, 7) # explain digit "6"
 heatmaps = heatmap(expl; process_batch=true)
-mosaic(heatmaps; nrow=10)
+mosaic(heatmaps; nrow=5)
 
 #md # !!! note "Output type consistency"
 #md #

src/bibliography.jl

@@ -1,4 +1,4 @@
 # Gradient methods:
 const REF_SMILKOV_SMOOTHGRAD = "Smilkov et al., *SmoothGrad: removing noise by adding noise*"
 const REF_SUNDARARAJAN_AXIOMATIC = "Sundararajan et al., *Axiomatic Attribution for Deep Networks*"
-const REF_SELVARAJU_GRADCAM = "Selvaraju et al., *Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization*"
+const REF_SELVARAJU_GRADCAM = "Selvaraju et al., *Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization*"

src/gradcam.jl

@@ -19,7 +19,7 @@ struct GradCAM{F,A} <: AbstractXAIMethod
     feature_layers::F
     adaptation_layers::A
 end
-function (analyzer::GradCAM)(input, ns::AbstractNeuronSelector)
+function (analyzer::GradCAM)(input, ns::AbstractOutputSelector)
     A = analyzer.feature_layers(input)  # feature map
     feature_map_size = size(A, 1) * size(A, 2)
 

src/gradient.jl

@@ -1,4 +1,4 @@
-function gradient_wrt_input(model, input, ns::AbstractNeuronSelector)
+function gradient_wrt_input(model, input, ns::AbstractOutputSelector)
     output, back = Zygote.pullback(model, input)
     output_indices = ns(output)
 
@@ -19,7 +19,7 @@ struct Gradient{M} <: AbstractXAIMethod
     Gradient(model) = new{typeof(model)}(model)
 end
 
-function (analyzer::Gradient)(input, ns::AbstractNeuronSelector)
+function (analyzer::Gradient)(input, ns::AbstractOutputSelector)
     grad, output, output_indices = gradient_wrt_input(analyzer.model, input, ns)
     return Explanation(grad, output, output_indices, :Gradient, :sensitivity, nothing)
 end
@@ -35,7 +35,7 @@ struct InputTimesGradient{M} <: AbstractXAIMethod
     InputTimesGradient(model) = new{typeof(model)}(model)
 end
 
-function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
+function (analyzer::InputTimesGradient)(input, ns::AbstractOutputSelector)
     grad, output, output_indices = gradient_wrt_input(analyzer.model, input, ns)
     attr = input .* grad
     return Explanation(

src/input_augmentation.jl

@@ -3,7 +3,7 @@
 
 Neuron selector that passes through an augmented neuron selection.
 """
-struct AugmentationSelector{I} <: AbstractNeuronSelector
+struct AugmentationSelector{I} <: AbstractOutputSelector
     indices::I
 end
 (s::AugmentationSelector)(out) = s.indices
@@ -103,7 +103,7 @@ function NoiseAugmentation(analyzer, n, σ::Real=0.1f0, args...)
     return NoiseAugmentation(analyzer, n, Normal(0.0f0, Float32(σ)^2), args...)
 end
 
-function (aug::NoiseAugmentation)(input, ns::AbstractNeuronSelector)
+function (aug::NoiseAugmentation)(input, ns::AbstractOutputSelector)
     # Regular forward pass of model
     output = aug.analyzer.model(input)
     output_indices = ns(output)
@@ -142,7 +142,7 @@ struct InterpolationAugmentation{A<:AbstractXAIMethod} <: AbstractXAIMethod
 end
 
 function (aug::InterpolationAugmentation)(
-    input, ns::AbstractNeuronSelector; input_ref=zero(input)
+    input, ns::AbstractOutputSelector; input_ref=zero(input)
 )
     size(input) != size(input_ref) &&
         throw(ArgumentError("Input reference size doesn't match input size."))

test/Project.toml

@@ -10,11 +10,3 @@ Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 XAIBase = "9b48221d-a747-4c1b-9860-46a1d8ba24a7"
 
-[compat]
-Aqua = "0.8"
-Distributions = "0.25"
-Flux = "0.13, 0.14"
-JLD2 = "0.4"
-ReferenceTests = "0.10"
-Suppressor = "0.2"
-XAIBase = "1.2"