using MLDatasets
using Images
using Augmentor
using Statistics
using Flux: onehotbatch

# Compute train dataset statistics
if !isdefined(Main, :meanC10) || !isdefined(Main, :stdC10)
    r = CIFAR10.traintensor()[:,:,1,:]
    g = CIFAR10.traintensor()[:,:,2,:]
    b = CIFAR10.traintensor()[:,:,3,:]
    const meanC10 = [mean(r), mean(g), mean(b)]
    const stdC10 = [std(r), std(g), std(b)]
end

function normalize(img)
    tmp = Float32.((channelview(img) .- meanC10) ./ stdC10)
    return permutedims(tmp, (3, 2, 1))
end