ImageMorphology.erode for a DiskArray

[bjarthur]

i was hoping a simple little rewrite of this section of code to use broadcasting would magically make erode fast for a boolean Zarr array, but alas, no:

diff --git a/src/extreme_filter.jl b/src/extreme_filter.jl
index 1d63ace..35af06c 100644
--- a/src/extreme_filter.jl
+++ b/src/extreme_filter.jl
@@ -527,10 +527,11 @@ function _extreme_filter_bool!(f, out, A::AbstractArray{Bool}, Ω)
     R_inner = (first(R) + δ):(last(R) - δ)
 
     select = _fast_select(f)
-    @inbounds for p in R_inner
-        out[p] = select(A, p, Ω)
-    end
+    @info 1
+    out[R_inner] .= select.(Ref(A), R_inner, Ref(Ω))
+    @info 2
     for p in EdgeIterator(R, R_inner)
+        @info p
         out[p] = select(A, p, Ω)
     end
     return out

my test harness prints "1" and then hangs for the call to erode! with the ZArray:

using ImageMorphology, Zarr
a = rand(Bool, 256,256,256);
z = ZArray(a);
oa = similar(a);
oz = similar(z);
se = strel_diamond((3, 3, 3));
erode!(oa, a, se);
erode!(oz, z, se);

i'm guessing it's because the logic in select has too many branches, or perhaps because it indexes the Zarray on two different lines (544 and 550).

is there any hope to an elegant solution here? if so, i'll keep banging my head on it.

in the meantime, i've got this hack:

using ImageMorphology

for (fun, filt) in (("erode!", typemax), ("dilate!", typemin))
    fun_sym = Symbol(fun)
    @eval ImageMorphology.$fun_sym(out, img; dims=coords_spatial(img), r=nothing, chunk) =
        ImageMorphology.$fun_sym(out, img, strel_box(img, dims; r), chunk=chunk)

    @eval @doc """
        $($fun)(out, img, se; chunk)
    
    $($fun) `img` in chunks.  Doing so can be much faster than the generic
    element-wise algorithm if `img` is backed by, for example, a Zarr array.
    `chunk` should be a tuple with an element for each dimension.  The chunk
    can nominally be of any size but generally should be larger than `se` and
    the underlying chunk size of `img`.

    # Examples

    $($fun)(out, img, se; chunk=(64,64,64))
    """ ImageMorphology.$fun_sym
    @eval function ImageMorphology.$fun_sym(out, img, se; chunk)
        pad = extrema.(axes(se))
        PAD1 = CartesianIndex(ntuple(i->pad[i][1], length(pad)))
        PAD2 = CartesianIndex(ntuple(i->pad[i][2], length(pad)))
        ZERO = CartesianIndex(ntuple(_->0, ndims(img)))
        ONE = CartesianIndex(ntuple(_->1, ndims(img)))
        SIZEIMG = CartesianIndex(size(img))
        CHUNK = CartesianIndex(chunk)
        chunk_pad = [c+length(range(p...))-1 for (c,p) in zip(chunk,pad)]
        img_chunk = Array{eltype(img)}(undef, chunk_pad...)
        out_chunk = similar(img_chunk)
        for I in CartesianIndices((:).(ntuple(_->1, ndims(img)), chunk, size(img)))
            img_chunk .= $filt(eltype(img))
            R = min(SIZEIMG, I+CHUNK-ONE)
            IMG1 = max(ONE, I+PAD1)
            IMG2 = min(SIZEIMG, R+PAD2)
            img_chunk[ONE:IMG2-IMG1+ONE] .= img[IMG1:IMG2]
            ImageMorphology.$fun_sym(out_chunk, img_chunk, se)
            OUT1 = ONE-PAD1 + min(ZERO, I+PAD1-ONE)
            OUT2 = OUT1+R-I
            out[I:R] .= out_chunk[OUT1:OUT2]
        end
    end
end

thanks in advance for any help!

[rafaqz]

It's not magic...

Rather than hoping for it not to hang, you can test it with DiskArrays.allowscalar(false). Your code will surely error due to scalar indexing, and that will be why it's slow.

DiskArrayEngine.jl kind of is magic and can probably do this (but not with a broadcast).

You can also just try DiskArrays.cache on whichever array is read with the erode window, and it might just work fine

[meggart]

Yes, this is a bit more tricky because you will need to apply the function to overlapping boxes of subsets of your data and you can not just go chunk by chunk. Indeed this is exactly a use case that DiskArrayEngine is designed for. The interface for DAE is still a bit clunky, but I will post a solution for your example here soon.

[meggart]

Ok, so here is an update how you could do this with DiskArrayEngine. I tried to make a function that would work for several functions from ImageMorphology. I am still not really happy because of the necessary allocation of the temp array.

using ImageMorphology, Zarr
import DiskArrayEngine as DAE


#This is a bit unlucky, we need to allocate an intermediate array that has the same size as the output
#This is because the output is smaller than the input and does not contain the "halo" of the input
function apply_inner(xout,xin,f,se,pad;dims=:,threaded=false)
    if size(xout) != size(xin)
        xout_ = similar(xin)
        f(xout_,xin,se)
        copyto!(xout,view(xout_,axes(xout)...))
    else
        f(xout,xin,se)
    end
end

function apply_imagemorphology_function(f,img,se;return_type=eltype(img))
    pad = extrema.(axes(se))
    widths = map(length, axes(se))
    windows = DAE.MovingWindow.(1 .+ first.(pad),1,widths,size(img),map(s->(1,s),size(img)))
    inars = DAE.InputArray.((img,),windows=windows);
    outars = DAE.create_outwindows(size(img));
    uf = DAE.create_userfunction(apply_inner,return_type,is_blockfunction=true,is_mutating=true,args=(f,se,pad))

    op = DAE.GMDWop(inars,(outars,),uf)
    DAE.results_as_diskarrays(op)[1]
end


a = rand(Bool, 256,256,256);
img = ZArray(a,chunks = (32,32,32));
se = strel_diamond((3, 3, 3));

#This is a lzay represntation of the result
res = apply_imagemorphology_function(erode!,img,se)
#But it can be indexed into like a normal array
res[90:100,90:100,100]

#The whole result can be computed into an array
aa = DAE.compute(res);

#Or into a preallocated array on disk
output = zzeros(Bool,size(img)...,chunks = (16,16,64),path=tempname())
aa = DAE.compute!(output,res);

#One can control how much the cache size and how much memor is used during the computation
# by providing the max_cache keyword (in bytes)
# The following splits up the computatio in much smaller while still trying to align with the chunks 
# of both input and output as good as possible
aa = DAE.compute!(output,res,max_cache=1e6);

in order to improve on this it would be great if erode! would work with arrays of different sizes where the output axes might only cover a subset of the inputs like this:

using OffsetArrays
img = OffsetArray(rand(Bool,54,54),-2,-2)
output = zeros(Bool,50,50)
erode!(output,img,strel_diamond((3,3)))

[bjarthur]

DAE seems like it could be very convenient, but in my hands the above example for erode is several fold slower for a 2k x 2k x 2k zarr than my hand-tuned code further above. perhaps because of the allocated intermediate array. am looking into that.

in the meantime, what is being optimized?

julia> DAE.compute!(scratch3,res)
┌ Warning: The selected optimization algorithm requires second order derivatives, but `SecondOrder` ADtype was not provided. 
│         So a `SecondOrder` with ADTypes.AutoForwardDiff() for both inner and outer will be created, this can be suboptimal and not work in some cases so 
│         an explicit `SecondOrder` ADtype is recommended.
└ @ OptimizationBase ~/.julia/packages/OptimizationBase/UXLhR/src/cache.jl:49

also, how do you use the cache? are you traversing the chunks in morton order so as to reduce cache misses?

would be great if DAE were documented. am curious how it works internally.

[felixcremer]

The optimization is in optimize_loopranges and the idea as I understand it, is that we would like to reduce the number of caches that we have to open multiple times. This optimization gets the chunking structure of the input arrays and the function that is applied on the inputs and the amount of memory that is available to determine, how many chunks can be loaded at the same time.