Add ImageNet preprocessing utilities (#80)

Author: adrhill

Project.toml

@@ -8,6 +8,7 @@ ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
+ImageTransformations = "02fcd773-0e25-5acc-982a-7f6622650795"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

src/ExplainableAI.jl

@@ -10,6 +10,7 @@ using Tullio
 
 # Heatmapping:
 using ImageCore
+using ImageTransformations: imresize
 using ColorSchemes
 
 # Model checks:
@@ -29,7 +30,7 @@ include("lrp_checks.jl")
 include("lrp_rules.jl")
 include("lrp.jl")
 include("heatmap.jl")
-
+include("imagenet.jl")
 export analyze
 
 # Analyzers
@@ -53,5 +54,5 @@ export heatmap
 
 # utils
 export strip_softmax, flatten_model, flatten_chain, canonize
-
+export preprocess_imagenet
 end # module

src/imagenet.jl

@@ -0,0 +1,32 @@
+# Image preprocessing for ImageNet models.
+# Code adapted from Metalhead 0.5.3's deprecated utils.jl
+# TODO: Remove once matching functionality is in either Metalhead.jl or MLDatasets.jl
+
+# Coefficients taken from PyTorch's ImageNet normalization code
+const PYTORCH_MEAN = [0.485f0, 0.456f0, 0.406f0]
+const PYTORCH_STD = [0.229f0, 0.224f0, 0.225f0]
+const IMGSIZE = (224, 224)
+
+# Take rectangle of pixels of shape `outsize` at the center of image `im`
+adjust(i::Integer) = ifelse(iszero(i % 2), 1, 0)
+function center_crop_view(im::AbstractMatrix, outsize=IMGSIZE)
+    im = imresize(im; ratio=maximum(outsize .// size(im)))
+    h2, w2 = div.(outsize, 2) # half height, half width of view
+    h_adjust, w_adjust = adjust.(outsize)
+    return @view im[
+        ((div(end, 2) - h2):(div(end, 2) + h2 - h_adjust)) .+ 1,
+        ((div(end, 2) - w2):(div(end, 2) + w2 - w_adjust)) .+ 1,
+    ]
+end
+
+"""
+    preprocess_imagenet(img)
+
+Preprocess an image for use with Metalhead.jl's ImageNet models using PyTorch weights.
+Uses PyTorch's normalization constants.
+"""
+function preprocess_imagenet(im::AbstractMatrix{<:AbstractRGB}, T=Float32::Type{<:Real})
+    im = center_crop_view(im)
+    im = (channelview(im) .- PYTORCH_MEAN) ./ PYTORCH_STD
+    return convert.(T, PermutedDimsArray(im, (3, 2, 1))) # Convert Image.jl's CHW to WHC
+end

test/references/utils/preprocess_imagnet.jld2

@@ -20,6 +20,10 @@ using ReferenceTests
         @info "Running tests on heatmaps..."
         include("test_heatmaps.jl")
     end
+    @testset "ImageNet preprocessing" begin
+        @info "Running tests on ImageNet preprocessing..."
+        include("test_imagenet.jl")
+    end
     @testset "Canonize" begin
         @info "Running tests on model canonization..."
         include("test_canonize.jl")

test/runtests.jl

@@ -0,0 +1,10 @@
+using ExplainableAI
+using ImageCore
+
+A = RGB{Float32}[
+    RGB{Float32}(0.44557732f0, 0.25328094f0, 0.53720146f0) RGB{Float32}(0.99433f0, 0.37066674f0, 0.8781263f0) RGB{Float32}(0.59815156f0, 0.21008879f0, 0.07259983f0)
+    RGB{Float32}(0.6966612f0, 0.27341717f0, 0.40360665f0) RGB{Float32}(0.12119287f0, 0.63196003f0, 0.32167268f0) RGB{Float32}(0.31825548f0, 0.7599565f0, 0.20566207f0)
+]
+x = preprocess_imagenet(A)
+@test_reference "references/utils/preprocess_imagnet.jld2" Dict("x" => x) by =
+    (r, a) -> isapprox(r["x"], a["x"]; rtol=0.05)