Support fillval in flood_fill! (#77)

This also distinguishes `flood` (which always returns a "mask") from
`flood_fill!` (which can be used to insert non-boolean values).

Also adds docstrings and doctests

Author: timholy

Project.toml

@@ -39,7 +39,8 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Documenter", "FileIO", "ImageMagick", "Images", "Test"]
+test = ["Documenter", "FileIO", "ImageMagick", "Images", "SparseArrays", "Test"]

src/ImageSegmentation.jl

@@ -34,7 +34,7 @@ export
     felzenszwalb,
     fast_scanning,
     fast_scanning!,
-    flood_fill,
+    flood,
     flood_fill!,
     watershed,
     hmin_transform,

src/flood_fill.jl

@@ -1,35 +1,128 @@
-flood_fill(f, src::AbstractArray, idx::Union{Integer,CartesianIndex}, nbrhood_function=diamond_iterator(window_neighbors(src))) =
+"""
+    mask = flood(f, src, idx, nbrhood_function=diamond_iterator((3,3,...)))
+
+Return an array `mask` with the same axes as `src`, marked `true` for all elements of `src` that:
+
+- satisfy `f(src[i]) == true` and
+- are connected by such elements to the starting point `idx` (an integer index or `CartesianIndex`).
+
+This throws an error if `f` evaluates as `false` for the starting value `src[idx]`.
+The sense of connectivity is defined by `nbrhood_function`, with two choices being
+[`ImageSegmentation.diamond_iterator`](@ref) and [`ImageSegmentation.box_iterator`](@ref.)
+
+# Examples
+
+```jldoctest; setup=:(using ImageSegmentation, ImageCore)
+julia> mostly_red(c) = red(c) > green(c) && red(c) > blue(c)
+mostly_red (generic function with 1 method)
+
+julia> img = repeat(LinRange(colorant"red", colorant"blue", 4), 1, 2) # red-to-blue
+4×2 Array{RGB{Float32},2} with eltype RGB{Float32}:
+ RGB{Float32}(1.0,0.0,0.0)            RGB{Float32}(1.0,0.0,0.0)
+ RGB{Float32}(0.666667,0.0,0.333333)  RGB{Float32}(0.666667,0.0,0.333333)
+ RGB{Float32}(0.333333,0.0,0.666667)  RGB{Float32}(0.333333,0.0,0.666667)
+ RGB{Float32}(0.0,0.0,1.0)            RGB{Float32}(0.0,0.0,1.0)
+
+julia> flood(mostly_red, [img; img], 1)   # only first copy of `img` is connected
+8×2 BitMatrix:
+ 1  1
+ 1  1
+ 0  0
+ 0  0
+ 0  0
+ 0  0
+ 0  0
+ 0  0
+```
+
+See also [`flood_fill!`](@ref).
+"""
+flood(f, src::AbstractArray, idx::Union{Integer,CartesianIndex}, nbrhood_function=diamond_iterator(window_neighbors(src))) =
     flood_fill!(f, falses(axes(src)) #=fill!(similar(src, Bool), false)=#, src, idx, nbrhood_function)
 
-function flood_fill(src::AbstractArray, idx::Union{Integer,CartesianIndex},
+function flood(src::AbstractArray, idx::Union{Integer,CartesianIndex},
                     nbrhood_function=diamond_iterator(window_neighbors(src)); thresh)
     validx = src[idx]
     validx = accum_type(typeof(validx))(validx)
     return let validx=validx
-        flood_fill(val -> default_diff_fn(val, validx) < thresh, src, idx, nbrhood_function)
+        flood(val -> default_diff_fn(val, validx) < thresh, src, idx, nbrhood_function)
     end
 end
 
-function flood_fill!(f, dest, src::AbstractArray, idx::Union{Int,CartesianIndex}, nbrhood_function=diamond_iterator(window_neighbors(src)))
+"""
+    flood_fill!(f, dest, src, idx, nbrhood_function=diamond_iterator((3,3,...)); fillvalue=true, isfilled = isequal(fillvalue))
+
+Set entries of `dest` to `fillvalue` for all elements of `src` that:
+
+- satisfy `f(src[i]) == true` and
+- are connected by such elements to the starting point `idx` (an integer index or `CartesianIndex`).
+
+This throws an error if `f` evaluates as `false` for the starting value `src[idx]`.
+The sense of connectivity is defined by `nbrhood_function`, with two choices being
+[`ImageSegmentation.diamond_iterator`](@ref) and [`ImageSegmentation.box_iterator`](@ref.)
+
+You can optionally omit `dest`, in which case entries in `src` will be set to `fillvalue`.
+You may also supply `isfilled`, which should return `true` for any value in `dest`
+which does not need to be set or visited; one requirement is that `isfilled(fillvalue) == true`.
+
+# Examples
+
+```jldoctest; setup=:(using ImageSegmentation)
+julia> a = repeat([1:4; 1:4], 1, 3)
+8×3 Matrix{Int64}:
+ 1  1  1
+ 2  2  2
+ 3  3  3
+ 4  4  4
+ 1  1  1
+ 2  2  2
+ 3  3  3
+ 4  4  4
+
+julia> flood_fill!(>=(3), a, CartesianIndex(3, 2); fillvalue = -1, isfilled = <(0))
+8×3 Matrix{Int64}:
+  1   1   1
+  2   2   2
+ -1  -1  -1
+ -1  -1  -1
+  1   1   1
+  2   2   2
+  3   3   3
+  4   4   4
+```
+
+See also [`flood`](@ref).
+"""
+function flood_fill!(f,
+                     dest,
+                     src::AbstractArray,
+                     idx::Union{Int,CartesianIndex},
+                     nbrhood_function = diamond_iterator(window_neighbors(src));
+                     fillvalue = true,
+                     isfilled = (fillvalue === true) ? identity : isequal(fillvalue))
     R = CartesianIndices(src)
-    axes(dest) == R.indices || throw(DimensionMismatch("$(axes(dest)) do not match $(Tuple(R))"))
     idx = R[idx]  # ensure cartesian indexing
     f(src[idx]) || throw(ArgumentError("starting point fails to meet criterion"))
     q = [idx]
-    _flood_fill!(f, dest, src, R, q, nbrhood_function)
+    fillvalue = convert(eltype(dest), fillvalue)
+    axes(dest) == R.indices || throw(DimensionMismatch("$(axes(dest)) do not match $(Tuple(R))"))
+    _flood_fill!(f, dest, src, R, q, nbrhood_function, fillvalue, isfilled)
     return dest
 end
-flood_fill!(f, dest, src::AbstractArray, idx::Integer, nbrhood_function=diamond_iterator(window_neighbors(src))) =
-    flood_fill!(f, dest, src, Int(idx)::Int, nbrhood_function)
+flood_fill!(f, dest, src::AbstractArray, idx::Integer, args...; kwargs...) =
+    flood_fill!(f, dest, src, Int(idx)::Int, args...; kwargs...)
+flood_fill!(f, src::AbstractArray, idx::Union{Integer,CartesianIndex}, args...; kwargs...) =
+    flood_fill!(f, src, src, idx, args...; kwargs...)
 
 # This is a trivial implementation (just to get something working), better would be a raster implementation
-function _flood_fill!(f::F, dest, src, R::CartesianIndices{N}, q, nbrhood_function::FN) where {F,N,FN}
+function _flood_fill!(f::F, dest, src, R::CartesianIndices{N}, q, nbrhood_function::FN, fillvalue, isfilled::C) where {F,N,FN,C}
+    isfilled(fillvalue) == true || throw(ArgumentError("`isfilled(fillvalue)` must return `true`"))
     while !isempty(q)
         idx = pop!(q)
-        dest[idx] = true
+        dest[idx] = fillvalue
         @inbounds for j in nbrhood_function(idx)
             j ∈ R || continue
-            if f(src[j]) && !dest[j]
+            if f(src[j]) && !isfilled(dest[j])
                 push!(q, j)
             end
         end

test/flood_fill.jl

@@ -3,19 +3,20 @@ using ImageSegmentation.Colors
 using ImageSegmentation.FixedPointNumbers
 using FileIO
 using Statistics
+using SparseArrays
 using Test
 
 @testset "flood_fill" begin
     # 0d
     a = reshape([true])
-    @test flood_fill(identity, a, CartesianIndex()) == a
-    @test_throws ArgumentError flood_fill(!, a, CartesianIndex())
+    @test flood(identity, a, CartesianIndex()) == a
+    @test_throws ArgumentError flood(!, a, CartesianIndex())
     # 1d
     a = 1:7
-    @test flood_fill(==(2), a, CartesianIndex(2)) == (a .== 2)
-    @test_throws ArgumentError flood_fill(==(2), a, CartesianIndex(3))
-    @test flood_fill(x -> 1 < x < 4, a, CartesianIndex(2)) == [false, true, true, false, false, false, false]
-    @test flood_fill(isinteger, a, CartesianIndex(2)) == trues(7)
+    @test flood(==(2), a, CartesianIndex(2)) == (a .== 2)
+    @test_throws ArgumentError flood(==(2), a, CartesianIndex(3))
+    @test flood(x -> 1 < x < 4, a, CartesianIndex(2)) == [false, true, true, false, false, false, false]
+    @test flood(isinteger, a, CartesianIndex(2)) == trues(7)
     # 2d
     ab = [true false false false;
          true true false false;
@@ -26,13 +27,13 @@ using Test
         for (f, a) in ((identity, ab), (==(1), an0f8), (==(1), agray))
         for idx in CartesianIndices(a)
             if f(a[idx])
-                @test flood_fill(f, a, idx) == a
+                @test flood(f, a, idx) == a
             else
-                @test_throws ArgumentError flood_fill(f, a, idx)
+                @test_throws ArgumentError flood(f, a, idx)
             end
         end
     end
-    @test flood_fill(identity, ab, Int16(1)) == ab
+    @test flood(identity, ab, Int16(1)) == ab
     # 3d
     k = 10
     a = falses(k, k, k)
@@ -47,19 +48,48 @@ using Test
     end
     for idx in eachindex(a)
         if a[idx]
-            @test flood_fill(identity, a, idx) == a
+            @test flood(identity, a, idx) == a
         else
-            @test_throws ArgumentError flood_fill(identity, a, idx)
+            @test_throws ArgumentError flood(identity, a, idx)
         end
     end
     # Colors
     path = download("https://github.com/JuliaImages/juliaimages.github.io/raw/source/docs/src/pkgs/segmentation/assets/flower.jpg")
     img = load(path)
-    seg = flood_fill(img, CartesianIndex(87,280); thresh=0.3)
+    seg = flood(img, CartesianIndex(87,280); thresh=0.3)
     @test 0.2*length(seg) <= sum(seg) <= 0.25*length(seg)
     c = mean(img[seg])
     # N0f8 makes for easier approximate testing
     @test N0f8(red(c)) ≈ N0f8(0.855)
     @test N0f8(green(c)) ≈ N0f8(0.161)
     @test N0f8(blue(c)) ≈ N0f8(0.439)
+
+    # flood_fill!
+    near3(x) = round(Int, x) == 3
+    a0 = [range(2, 4, length=9);]
+    a = copy(a0)
+    idx = (length(a)+1)÷2
+    dest = fill!(similar(a, Bool), false)
+    @test flood_fill!(near3, dest, a, idx) == (round.(a) .== 3)
+    a = copy(a0)
+    flood_fill!(near3, a, idx; fillvalue=3)
+    @test a == [near3(a0[i]) ? 3 : a[i] for i in eachindex(a)]
+    a = copy(a0)
+    flood_fill!(near3, a, idx; fillvalue=-1)
+    @test a == [near3(a0[i]) ? -1 : a[i] for i in eachindex(a)]
+    a = copy(a0)
+    @test_throws ArgumentError flood_fill!(near3, a, idx; fillvalue=-1, isfilled=near3)
+
+    # This mimics a "big data" application in which we have several structures we want
+    # to label with different segment numbers, and the `src` array is too big to fit
+    # in memory.
+    # It would be better to use a package like SparseArrayKit, which allows efficient
+    # insertions and supports arbitrary dimensions.
+    a = Bool[0 0 0 0 0 0 1 1;
+             1 1 0 0 0 0 0 0]
+    dest = spzeros(Int, size(a)...)   # stores the nonzero indexes in a Dict
+    flood_fill!(identity, dest, a, CartesianIndex(2, 1); fillvalue=1)
+    flood_fill!(identity, dest, a, CartesianIndex(1, 7); fillvalue=2)
+    @test dest == [0 0 0 0 0 0 2 2;
+                   1 1 0 0 0 0 0 0]
 end