Make ImageToTensor compatible with Flux's WHC format

src/preprocessing.jl

@@ -133,8 +133,8 @@ end
 """
     ImageToTensor()
 
-Expands an `Image{N, T}` of size `sz` to an `ArrayItem{N+1}` with
-size `(sz..., ch)` where `ch` is the number of color channels of `T`.
+Expands an `Image{N, T}` of size `(height, width, ...)` to an `ArrayItem{N+1}` with
+size `(width, height, ..., ch)` where `ch` is the number of color channels of `T`.
 
 Supports `apply!`.
 
@@ -144,9 +144,10 @@ Supports `apply!`.
 ```julia
 using DataAugmentation, Images
 
-image = Image(rand(RGB, 50, 50))
+h, w = 40, 50
+image = Image(rand(RGB, h, w))
 tfm = ImageToTensor()
-apply(tfm, image)
+apply(tfm, image) # ArrayItem in WHC format of size (50, 40, 3)
 ```
 
 """
@@ -165,24 +166,12 @@ function apply!(buf, ::ImageToTensor, image::Image; randstate = nothing)
     return buf
 end
 
-function imagetotensor(image::AbstractArray{C, N}, T = Float32) where {C<:Color, N}
-    T.(PermutedDimsArray(_channelview(image), ((i for i in 2:N+1)..., 1)))
+function imagetotensor(image::AbstractArray{C, N}, T = Float32) where {C<:Colorant, N}
+    T.(PermutedDimsArray(_channelview(image), (3, 2, 4:N+1..., 1)))
 end
 
-#=
-function imagetotensor(image::AbstractArray{C, N}, T = Float32) where {TC, C<:Color{TC, 1}, N}
-    return T.(_channelview(image))
-end
-=#
-
-
-# TODO: relax color type constraint, implement for other colors
-# single-channel colors need a `channelview` that also expands the array
-function imagetotensor!(buf, image::AbstractArray{<:Color, N}) where N
-    permutedims!(
-        buf,
-        _channelview(image),
-        (2:N+1..., 1))
+function imagetotensor!(buf, image::AbstractArray{<:Colorant, N}) where N
+    permutedims!(buf, _channelview(image), (3, 2, 4:N+1..., 1))
 end
 
 function tensortoimage(a::AbstractArray)
@@ -197,22 +186,20 @@ function tensortoimage(a::AbstractArray)
     end
 end
 
-function tensortoimage(C::Type{<:Color}, a::AbstractArray{T, N}) where {T, N}
-    perm = (N, 1:N-1...)
+function tensortoimage(C::Type{<:Colorant}, a::AbstractArray{T, N}) where {T, N}
+    perm = (N, 2, 1, 3:N-1...)
     return _colorview(C, PermutedDimsArray(a, perm))
 end
 
-
-function _channelview(img)
-    chview = channelview(img)
-    # for single-channel colors, expand the color dimension anyway
-    if size(img) == size(chview)
-        chview = reshape(chview, 1, size(chview)...)
-    end
-    return chview
+# For single-channel colors, expand the color dimension anyway
+# such that the output is always of size (channels, height, width, ...)
+_channelview(img::AbstractArray{<:Colorant{T, N}}) where {T, N} = channelview(img)
+function _channelview(img::AbstractArray{<:Colorant{T, 1}}) where T
+    cv = channelview(img)
+    return reshape(cv, 1, size(cv)...)
 end
 
-function _colorview(C::Type{<:Color}, img)
+function _colorview(C::Type{<:Colorant}, img)
     if size(img, 1) == 1
         img = reshape(img, size(img)[2:end])
     end

test/preprocessing.jl

@@ -102,7 +102,7 @@ end
 
     res = apply(tfm, item1)
     a = itemdata(res)
-    @test size(a) == (32, 48, 3)
+    @test size(a) == (48, 32, 3)
     @test eltype(a) == Float32
 
     testapply(tfm, item1)