Flux implementation of UNet

unet.jl

function (u::UNet)(x::AbstractArray)
    enc_out = []

    out = x
    for i in 1:4
        out = u.encoder[i](out)
        push!(enc_out, out)

        out = u.pool[i](out)
    end

    out = u.bottleneck(out)

    for i in 4:-1:1
        out = u.upconv[5 - i](out)
        out = cat(out, enc_out[i]; dims = 3)
        out = u.decoder[i](out)
    end

    return σ(u.final_conv(out))
end