Merge pull request #24 from lorenzoh/OneHot-MaskMulti

Add one hot encoding for `MaskMulti`s

Author: lorenzoh

docs/literate/preprocessing.md

@@ -5,4 +5,5 @@ This library also implements some general transformations useful for getting dat
 - [`ToEltype`](#)`(T)` converts the element type of any [`AbstractArrayItem`](#) to `T`.
 - [`ImageToTensor`](#) converts an image to an `ArrayItem` with another dimension for the color channels 
 - [`Normalize`](#) normalizes image tensors
+- [`OneHot`](#) to one-hot encode multi-class masks ([`MaskMulti`](#)s)
 

src/DataAugmentation.jl

@@ -61,6 +61,7 @@ export Item,
     CenterResizeCrop,
     Buffered,
     BufferedThreadsafe,
+    OneHot,
     apply,
     Reflect,
     FlipX,

src/preprocessing.jl

@@ -114,7 +114,7 @@ denormalize(a, means, stds) = denormalize!(copy(a), means, stds)
 """
     NormalizeIntensity()
 
-Normalizes the pixels of an array based on calculated mean and std. 
+Normalizes the pixels of an array based on calculated mean and std.
 """
 
 struct NormalizeIntensity <: Transform end
@@ -183,6 +183,49 @@ end
 tensortoimage(a::AbstractArray{T, 3}) where T = colorview(RGB, permuteddimsview(a, (3, 1, 2)))
 tensortoimage(a::AbstractArray{T, 2}) where T = colorview(Gray, a)
 
+
+# OneHot encoding
+
+"""
+    OneHot([T = Float32])
+
+One-hot encodes a `MaskMulti` with `n` classes and size `sz` into
+an array item of size `(sz..., n)` with element type `T`. Supports [`apply!`].
+
+```julia
+item = MaskMulti(rand(1:4, 100, 100), 1:4)
+apply(OneHot(), item)
+```
+"""
+struct OneHot{T} <: Transform end
+OneHot() = OneHot{Float32}()
+
+function apply(tfm::OneHot{T}, item::MaskMulti; randstate = nothing) where T
+    mask = itemdata(item)
+    a = zeros(T, size(mask)..., length(item.classes))
+    for I in CartesianIndices(mask)
+        a[I, mask[I]] = one(T)
+    end
+
+    return ArrayItem(a)
+end
+
+
+function apply!(buf, tfm::OneHot{T}, item::MaskMulti; randstate = nothing) where T
+    mask = itemdata(item)
+    a = itemdata(buf)
+    @show a[1:6]
+    fill!(a, zero(T))
+    @show a[1:6]
+
+    for I in CartesianIndices(mask)
+        a[I, mask[I]] = one(T)
+    end
+
+    return buf
+end
+
+
 function onehot(T, x::Int, n::Int)
     v = fill(zero(T), n)
     v[x] = one(T)

test/preprocessing.jl

@@ -55,6 +55,22 @@ end
     end
 end
 
+@testset ExtendedTestSet "OneHot" begin
+    tfm = OneHot()
+    mask = rand(1:4, 10, 10)
+    item = MaskMulti(mask, 1:4)
+    @test_nowarn apply(tfm, item)
+    aitem = apply(tfm, item)
+    @test size(itemdata(aitem)) == (10, 10, 4)
+
+    item2 = MaskMulti(rand(1:3, 10, 10), 1:4)
+    buf = itemdata(aitem)
+    bufcopy = copy(buf)
+    apply!(aitem, tfm, item2)
+    @test itemdata(item) == itemdata(item2) || itemdata(aitem) != bufcopy
+
+end
+
 @testset ExtendedTestSet "Image pipeline" begin
     image = Image(rand(RGB, 150, 150))