Accepted mask type for Augmentor.jl

Cheers,

I’m a new Augmentor.jl user. Need to apply same augmentation to images and masks for training a semantic segmentation model. Augmentation of images has worked out smoothly. Issue is when augmenting masks as the process has refused both,an Int(H,W,N) array (each pixel labelled as an integer class), and its equivalent onehot Bool(H,W,C,N) array.

Code for RGB images and onehot masks as follows:

dims = (512,512)
pl = FlipX(0.5) |> Rotate([-30,-20,-10,0,10,20,30]) |> ColorJitter() |> CropSize(dims)

Xtrain # same type as Xaug, preloaded
ytrain # same type as yaug, preloaded
Xaug = Array{RGB,3}(undef, (dims...,Ntrain))
yaug = Array{Bool,4}(undef, (dims...,C,Ntrain))

augmentbatch!((Xaug), (Xtrain), pl)
augmentbatch!((yaug), Mask(ytrain), pl)

Remark: images and masks are augmented separately for debug purposes.

For this case, error message is as follows:

ERROR: MethodError: no method matching augmentbatch!(::Array{Bool, 4}, ::Mask{Array{…}}, ::Augmentor.ImmutablePipeline{4, Tuple{…}})
The function `augmentbatch!` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  augmentbatch!(::Union{Tuple, AbstractArray}, ::Union{Tuple, AbstractArray}, ::Any, Any...)
   @ Augmentor ~/.julia/packages/Augmentor/96doC/src/augmentbatch.jl:80
  augmentbatch!(::CPUThreads, ::AbstractVector{<:Union{Tuple, AbstractArray}}, ::AbstractVector{<:Union{Tuple, AbstractArray}}, ::Any)
   @ Augmentor ~/.julia/packages/Augmentor/96doC/src/augmentbatch.jl:110
  augmentbatch!(::CPU1, ::AbstractVector{<:Union{Tuple, AbstractArray}}, ::AbstractVector{<:Union{Tuple, AbstractArray}}, ::Any)
   @ Augmentor ~/.julia/packages/Augmentor/96doC/src/augmentbatch.jl:98
  ...

Thanks for advising.