augment! and augmentbatch!

NEWS.md

@@ -1,5 +1,12 @@
 # v0.2.0
 
+New functionality:
+
+- `augment!`: Use preallocated storage for the output
+
+- `augmentbatch!`: Augment a whole batch of images. Optionally
+  using multiple threads.
+
 New operations:
 
 - `ConvertEltype`: Convert the array elements to the given type

REQUIRE

@@ -10,6 +10,8 @@ StaticArrays
 OffsetArrays
 IdentityRanges
 ColorTypes 0.4
+MLDataPattern 0.1.2
+ComputationalResources 0.0.2
 ShowItLikeYouBuildIt
 FileIO
 Compat 0.17

src/Augmentor.jl

@@ -13,6 +13,8 @@ using Interpolations
 using StaticArrays
 using OffsetArrays
 using IdentityRanges
+using MLDataPattern
+using ComputationalResources
 using FileIO
 using ShowItLikeYouBuildIt
 using Compat
@@ -59,6 +61,8 @@ export
     ElasticDistortion,
 
     augment,
+    augment!,
+    augmentbatch!,
 
     testpattern
 
@@ -87,5 +91,10 @@ include("operations/distortion.jl")
 include("pipeline.jl")
 include("codegen.jl")
 include("augment.jl")
+include("augmentbatch.jl")
+
+function __init__()
+    rand_mutex[] = Threads.Mutex()
+end
 
 end # module

src/augment.jl

@@ -39,12 +39,46 @@ function augment(op::Union{AbstractPipeline,Operation})
     augment(use_testpattern(), op)
 end
 
-# --------------------------------------------------------------------
-
 @inline function _augment(img, pipeline::AbstractPipeline)
     _augment(img, operations(pipeline)...)
 end
 
 @generated function _augment(img, pipeline::Vararg{Operation})
-    Expr(:block, Expr(:meta, :inline), augment_impl(:img, pipeline))
+    Expr(:block, Expr(:meta, :inline), augment_impl(:img, pipeline, false))
+end
+
+# --------------------------------------------------------------------
+
+"""
+    augment!(out, img, pipeline) -> out
+
+Apply the operations of the given `pipeline` to the image `img`
+and write the resulting image into `out`.
+
+The parameter `pipeline` can be a subtype of
+`Augmentor.Pipeline`, a tuple of `Augmentor.Operation`, or a
+single `Augmentor.Operation`
+
+```julia
+img = testpattern()
+out = similar(img)
+augment!(out, img, FlipX() |> FlipY())
+augment!(out, img, (FlipX(), FlipY()))
+augment!(out, img, FlipX())
+```
+"""
+augment!(out, img, op::Operation) = augment!(out, img, (op,))
+
+function augment!(out, img, pipeline::AbstractPipeline)
+    out_lazy = _augment_avoid_eager(img, pipeline)
+    copy!(match_idx(out, indices(out_lazy)), out_lazy)
+    out
+end
+
+@inline function _augment_avoid_eager(img, pipeline::AbstractPipeline)
+    _augment_avoid_eager(img, operations(pipeline)...)
+end
+
+@generated function _augment_avoid_eager(img, pipeline::Vararg{Operation})
+    Expr(:block, Expr(:meta, :inline), augment_impl(:img, pipeline, true))
 end

src/augmentbatch.jl

@@ -0,0 +1,68 @@
+_berror() = throw(ArgumentError("Number of output images must be equal to the number of input images"))
+
+imagesvector(imgs::AbstractArray) = obsview(imgs)
+@inline imagesvector(imgs::AbstractVector{<:AbstractArray}) = imgs
+
+# --------------------------------------------------------------------
+
+"""
+    augmentbatch!([resource], outs, imgs, pipeline) -> outs
+
+Apply the operations of the given `pipeline` to the images in
+`imgs` and write the resulting images into `outs`.
+
+Both `outs` and `imgs` have to contain the same number of images.
+Each of the two variables can either be in the form of a higher
+dimensional array for which the last dimension enumerates the
+individual images, or alternatively in the form of a vector of
+arrays, for which each vector element denotes an image.
+
+The parameter `pipeline` can be a subtype of
+`Augmentor.Pipeline`, a tuple of `Augmentor.Operation`, or a
+single `Augmentor.Operation`.
+
+The optional first parameter `resource` can either be `CPU1()`
+(default) or `CPUThreads()`. In the case of the later the images
+will be augmented in parallel. For this to make sense make sure
+that the environment variable `JULIA_NUM_THREADS` is set to a
+reasonable number so that `Threads.nthreads()` is greater than 1.
+"""
+function augmentbatch!(
+        outs::AbstractArray,
+        imgs::AbstractArray,
+        pipeline)
+    augmentbatch!(CPU1(), outs, imgs, pipeline)
+end
+
+function augmentbatch!(
+        r::AbstractResource,
+        outs::AbstractArray,
+        imgs::AbstractArray,
+        pipeline)
+    augmentbatch!(r, imagesvector(outs), imagesvector(imgs), pipeline)
+    outs
+end
+
+function augmentbatch!(
+        ::CPU1,
+        outs::AbstractVector{<:AbstractArray},
+        imgs::AbstractVector{<:AbstractArray},
+        pipeline)
+    length(outs) == length(imgs) || _berror()
+    for i in 1:length(outs)
+        augment!(outs[i], imgs[i], pipeline)
+    end
+    outs
+end
+
+function augmentbatch!(
+        ::CPUThreads,
+        outs::AbstractVector{<:AbstractArray},
+        imgs::AbstractVector{<:AbstractArray},
+        pipeline)
+    length(outs) == length(imgs) || _berror()
+    Threads.@threads for i in 1:length(outs)
+        augment!(outs[i], imgs[i], pipeline)
+    end
+    outs
+end

src/codegen.jl

@@ -1,9 +1,11 @@
 """
-    seek_connected(f, N::Int, head::DataType, tail::Tuple) -> (N, seq)
+    seek_connected(f, N::Int, head::DataType, tail::Tuple) -> (N, remainder)
 
 Recursively scan a tuple of `DataType` (split into its `head` and
-`tail`) to compute the uninterrupted sequence `seq` of adjacent
-operations (and its length `N`) where the predicate `f` is true.
+`tail`) to compute the length `N` of the uninterrupted sequence
+of adjacent operations where the predicate `f` is true.
+Additionally the `remainder` of the tuple (without that sequence)
+is also returned.
 """
 @inline function seek_connected(f, N::Int, head::Type{<:Operation}, tail::Tuple)
     if f(head)
@@ -32,51 +34,56 @@ end
 
 # --------------------------------------------------------------------
 
-function augment_impl(var_offset::Int, op_offset::Int, pipeline::Tuple)
-    augment_impl(var_offset, op_offset, first(pipeline), Base.tail(pipeline))
+function augment_impl(var_offset::Int, op_offset::Int, pipeline::Tuple, args...)
+    augment_impl(var_offset, op_offset, first(pipeline), Base.tail(pipeline), args...)
 end
 
-function augment_impl(var_offset::Int, op_offset::Int, pipeline::Tuple{})
+function augment_impl(var_offset::Int, op_offset::Int, pipeline::Tuple{}, args...)
     :($(Symbol(:img_, var_offset)))
 end
 
-function augment_impl(var_offset::Int, op_offset::Int, head, tail::NTuple{N,DataType}) where N
+function augment_impl(var_offset::Int, op_offset::Int, head::DataType, tail::NTuple{N,DataType}, avoid_eager = false) where N
     var_in  = Symbol(:img_, var_offset)
     var_out = Symbol(:img_, var_offset+1)
     if supports_lazy(head, tail)
-        num_affine, rest_affine = uses_affinemap(head, tail) ? seek_connected(uses_affinemap, 0, head, tail) : (0, nothing)
+        # If reached there are at least two adjacent lazy operations
+        num_affine, after_affine = uses_affinemap(head, tail) ? seek_connected(uses_affinemap, 0, head, tail) : (0, nothing)
         num_special, _ = seek_connected(x->(supports_permute(x)||supports_view(x)||supports_stepview(x)), 0, head, tail)
-        num_lazy, rest_lazy = seek_connected(supports_lazy, 0, head, tail)
+        num_lazy, after_lazy = seek_connected(supports_lazy, 0, head, tail)
         if num_special >= num_affine
             quote
                 $var_out = unroll_applylazy($(Expr(:tuple, (:(pipeline[$i]) for i in op_offset:op_offset+num_lazy-1)...)), $var_in)
-                $(augment_impl(var_offset+1, op_offset+num_lazy, rest_lazy))
+                $(augment_impl(var_offset+1, op_offset+num_lazy, after_lazy, avoid_eager))
             end
         else
             quote
                 $var_out = unroll_applyaffine($(Expr(:tuple, (:(pipeline[$i]) for i in op_offset:op_offset+num_affine-1)...)), $var_in)
-                $(augment_impl(var_offset+1, op_offset+num_affine, rest_affine))
+                $(augment_impl(var_offset+1, op_offset+num_affine, after_affine, avoid_eager))
             end
         end
     else
-        if length(tail) == 0 || supports_eager(head) || !supports_lazy(head)
+        # At most "head" is lazy (i.e. tail[1] is surely not).
+        # Unless "avoid_eager==true" we prefer using "applyeager" in
+        # this case because there is no neighbour synergy and we
+        # assume "applyeager" is more efficient.
+        if !supports_lazy(head) || (!avoid_eager && (length(tail) == 0 || supports_eager(head)))
             quote
                 $var_out = applyeager(pipeline[$op_offset], $var_in)
-                $(augment_impl(var_offset+1, op_offset+1, tail))
+                $(augment_impl(var_offset+1, op_offset+1, tail, avoid_eager))
             end
-        else # use lazy because there is no special eager implementation
+        else
             quote
-                $var_out = unroll_applylazy(pipeline[$op_offset], $var_in)
-                $(augment_impl(var_offset+1, op_offset+1, tail))
+                $var_out = applylazy(pipeline[$op_offset], $var_in)
+                $(augment_impl(var_offset+1, op_offset+1, tail, avoid_eager))
             end
         end
     end
 end
 
-function augment_impl(varname, pipeline::NTuple{N,DataType}) where N
+function augment_impl(varname::Symbol, pipeline::NTuple{N,DataType}, args...) where N
     quote
         img_1 = $varname
-        $(augment_impl(1, 1, pipeline))
+        $(augment_impl(1, 1, pipeline, args...))
     end
 end
 
@@ -102,8 +109,8 @@ end
 #       end
 #   end
 
-function augment_impl(pipeline::AbstractPipeline)
-    augment_impl(:input_image, map(typeof, operations(pipeline)))
+function augment_impl(pipeline::AbstractPipeline; avoid_eager = false)
+    augment_impl(:input_image, map(typeof, operations(pipeline)), avoid_eager)
 end
 
-augment_impl(op::Operation) = augment_impl((op,))
+augment_impl(op::Operation; kw...) = augment_impl((op,); kw...)

src/distortionfields.jl

@@ -18,10 +18,10 @@ end
 function uniform_field(gridheight::Int, gridwidth::Int, scale, border, normalize)
     A = if !border
         @assert gridwidth > 2 && gridheight > 2
-        _2dborder!(rand(2, gridheight, gridwidth), .5)
+        _2dborder!(safe_rand(2, gridheight, gridwidth), .5)
     else
         @assert gridwidth > 0 && gridheight > 0
-        rand(2, gridheight, gridwidth)
+        safe_rand(2, gridheight, gridwidth)
     end::Array{Float64,3}
     broadcast!(*, A, A, 2.)
     broadcast!(-, A, A, 1.)

src/operations/cache.jl

@@ -86,10 +86,6 @@ CacheImage(buffer::AbstractArray) = CacheImageInto(buffer)
 
 @inline supports_lazy(::Type{<:CacheImageInto}) = true
 
-@inline match_idx(buffer::AbstractArray, inds::Tuple) = buffer
-@inline match_idx(buffer::Array, inds::NTuple{N,UnitRange}) where {N} =
-    OffsetArray(buffer, inds)
-
 applyeager(op::CacheImageInto, img) = applylazy(op, img)
 
 function applylazy(op::CacheImageInto, img)

src/operations/crop.jl

@@ -480,12 +480,12 @@ function rcropratio_indices(op::RCropRatio, img::AbstractMatrix)
     elseif nw < w
         x_max = w - nw + 1
         @assert x_max > 0
-        x = rand(1:x_max)
+        x = safe_rand(1:x_max)
         1:h, x:(x+nw-1)
     elseif nh < h
         y_max = h - nh + 1
         @assert y_max > 0
-        y = rand(1:y_max)
+        y = safe_rand(1:y_max)
         y:(y+nh-1), 1:w
     else
         error("unreachable code reached")

src/operations/either.jl

@@ -158,7 +158,7 @@ end
 
 function toaffinemap(op::Either, img)
     supports_affine(typeof(op)) || throw(MethodError(toaffinemap, (op, img)))
-    p = rand()
+    p = safe_rand()
     for (i, p_i) in enumerate(op.cum_chances)
         if p <= p_i
             return toaffinemap_common(op.operations[i], img)
@@ -179,7 +179,7 @@ for KIND in (:eager, :permute, :view, :stepview, :affine, :affineview)
     APP = startswith(String(KIND),"affine") ? Symbol(FUN, :_common) : FUN
     @eval function ($FUN)(op::Either, img)
         ($SUP)(typeof(op)) || throw(MethodError($FUN, (op, img)))
-        p = rand()
+        p = safe_rand()
         for (i, p_i) in enumerate(op.cum_chances)
             if p <= p_i
                 return ($APP)(op.operations[i], img)

src/operations/rotation.jl

@@ -69,13 +69,13 @@ function applypermute(::Rotate90, img::AbstractMatrix{T}) where T
     view(perm_img, reverse(idx[2]), idx[1])
 end
 
-function applypermute(::Rotate90, sub::SubArray{T,2,IT}) where {T,IT<:PermutedDimsArray{T,2,(2,1)}}
+function applypermute(::Rotate90, sub::SubArray{T,2,IT,<:NTuple{2,Range}}) where {T,IT<:PermutedDimsArray{T,2,(2,1)}}
     idx = map(StepRange, sub.indexes)
     img = parent(parent(sub))
     view(img, reverse(idx[2]), idx[1])
 end
 
-function applypermute(::Rotate90, sub::SubArray{T,2}) where T
+function applypermute(::Rotate90, sub::SubArray{T,2,IT,<:NTuple{2,Range}}) where {T,IT}
     idx = map(StepRange, sub.indexes)
     img = parent(sub)
     perm_img = PermutedDimsArray{T,2,(2,1),(2,1),typeof(img)}(img)
@@ -223,13 +223,13 @@ function applypermute(::Rotate270, img::AbstractMatrix{T}) where T
     view(perm_img, idx[2], reverse(idx[1]))
 end
 
-function applypermute(::Rotate270, sub::SubArray{T,2,IT}) where {T,IT<:PermutedDimsArray{T,2,(2,1)}}
+function applypermute(::Rotate270, sub::SubArray{T,2,IT,<:NTuple{2,Range}}) where {T,IT<:PermutedDimsArray{T,2,(2,1)}}
     idx = map(StepRange, sub.indexes)
     img = parent(parent(sub))
     view(img, idx[2], reverse(idx[1]))
 end
 
-function applypermute(::Rotate270, sub::SubArray{T,2}) where T
+function applypermute(::Rotate270, sub::SubArray{T,2,IT,<:NTuple{2,Range}}) where {T,IT}
     idx = map(StepRange, sub.indexes)
     img = parent(sub)
     perm_img = PermutedDimsArray{T,2,(2,1),(2,1),typeof(img)}(img)
@@ -321,7 +321,7 @@ Rotate(degree::Real) = Rotate(degree:degree)
 @inline supports_eager(::Type{<:Rotate}) = false
 
 function toaffinemap(op::Rotate, img::AbstractMatrix)
-    recenter(RotMatrix(deg2rad(Float64(rand(op.degree)))), center(img))
+    recenter(RotMatrix(deg2rad(Float64(safe_rand(op.degree)))), center(img))
 end
 
 function Base.show(io::IO, op::Rotate)

src/operations/scale.jl

@@ -74,19 +74,19 @@ Scale() = throw(MethodError(Scale, ()))
 Scale(::Tuple{}) = throw(MethodError(Scale, ((),)))
 Scale(factors...) = Scale(factors)
 Scale(factor::Union{AbstractVector,Real}) = Scale((factor, factor))
-Scale(factors::NTuple{N,Any}) where {N} = Scale(map(_vectorize, factors))
+Scale(factors::NTuple{N,Any}) where {N} = Scale(map(vectorize, factors))
 Scale(factors::NTuple{N,Range}) where {N} = Scale{N}(promote(factors...))
 function Scale(factors::NTuple{N,AbstractVector}) where N
     Scale{N}(map(Vector{Float64}, factors))
 end
 function (::Type{Scale{N}})(factors::NTuple{N,Any}) where N
-    Scale(map(_vectorize, factors))
+    Scale(map(vectorize, factors))
 end
 
 @inline supports_eager(::Type{<:Scale}) = false
 
 function toaffinemap(op::Scale{2}, img::AbstractMatrix)
-    idx = rand(1:length(op.factors[1]))
+    idx = safe_rand(1:length(op.factors[1]))
     @inbounds tfm = recenter(@SMatrix([Float64(op.factors[1][idx]) 0.; 0. Float64(op.factors[2][idx])]), center(img))
     tfm
 end