feature_transform: add multithreading

If Julia is being run with multiple threads, this now will default
to using all threads in `feature_transform`. The multithreaded
implementation is relatively straightforward, although the recursive algorithm
makes it a little hard to wrap your head around:

- `computeft!` for dimension `d` is dependent only on slices of dimensions `1:d`,
  and indeed for all except the final call to `voronoift!` just on dimensions `1:d-1`.
  Hence you can divide the last dimension `N` into chunks and give each to a thread.

- For the final `voronoift!` call along dimension `N`, it's a
  one-dimensional operation along this dimension. Hence you can just
  split the next-to-last dimension into chunks instead.

This change was responsible for performance improvements in `distance_transform`
noted in JuliaImages/image_benchmarks#1.

Julia 1.3 is required for `@spawn`

Author: timholy

.github/workflows/UnitTest.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
       matrix:
-        julia-version: ['1.0', '1', 'nightly']
+        julia-version: ['1.3', '1', 'nightly']
         os: [ubuntu-latest]
         arch: [x64]
         include:

Project.toml

@@ -12,8 +12,8 @@ TiledIteration = "06e1c1a7-607b-532d-9fad-de7d9aa2abac"
 Documenter = "0.24, 0.25"
 ImageCore = "0.9"
 Requires = "1"
-TiledIteration = "0.2, 0.3"
-julia = "1"
+TiledIteration = "0.3.1"
+julia = "1.3"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

src/ImageMorphology.jl

@@ -3,8 +3,7 @@ module ImageMorphology
 using ImageCore
 using ImageCore: GenericGrayImage
 using LinearAlgebra
-using Base.Cartesian # TODO: delete this
-using TiledIteration: EdgeIterator
+using TiledIteration: EdgeIterator, SplitAxis, SplitAxes
 using Requires
 
 include("convexhull.jl")

src/bwdist.jl

@@ -1,9 +1,13 @@
 module FeatureTransform
 
+using ..ImageMorphology: SplitAxis, SplitAxes
+using ..ImageCore
+
 export feature_transform, distance_transform
 
 """
-    feature_transform(I::AbstractArray{Bool, N}, [w=nothing]) -> F
+    feature_transform(img::AbstractArray{Bool, N};
+                      weights=nothing, nthreads=Threads.nthreads()) -> F
 
 Compute the feature transform of a binary image `I`, finding the
 closest "feature" (positions where `I` is `true`) for each location in
@@ -26,20 +30,45 @@ See also: [`distance_transform`](@ref).
 Transforms of Binary Images in Arbitrary Dimensions' [Maurer et al.,
 2003] (DOI: 10.1109/TPAMI.2003.1177156)
 """
-function feature_transform(I::AbstractArray{Bool,N}, w::Union{Nothing,NTuple{N}}=nothing) where N
+function feature_transform(img::AbstractArray{<:Union{Bool,AbstractGray{Bool}},N};
+                           weights::Union{Nothing,NTuple{N}}=nothing,
+                           nthreads::Int = length(img) < 1000 ? 1 : Threads.nthreads()) where N
+    nthreads > 0 || error("the number of threads must be positive")
+    N == 0 && return reshape([CartesianIndex()])
+    # Allocate the output
+    F = similar(img, CartesianIndex{N})
+    axsimg = axes(img)
     # To allocate temporary storage for voronoift!, compute one
     # element (so we have the proper type)
-    fi = first(CartesianIndices(axes(I)))
-    drft = DistRFT(fi, w, (), Base.tail(fi.I))
-    tmp = Vector{typeof(drft)}()
-
-    # Allocate the output
-    F = similar(I, CartesianIndex{N})
-
-    # Compute the feature transform (recursive algorithm)
-    computeft!(F, I, CartesianIndex(()), w, tmp)
-
-    F
+    fi = first(CartesianIndices(axsimg))
+    drft = DistRFT(fi, weights, (), Base.tail(fi.I))
+    if nthreads == 1
+        tmp = typeof(drft)[]
+        computeft!(F, img, axsimg, CartesianIndex(()), weights, tmp)
+        # Finish the last dimension (for threading, we avoid doing it in computeft!)
+        finishft!(F, img, axsimg, CartesianIndex(()), weights, tmp)
+    else
+        tasks = Task[]
+        tmps = [typeof(drft)[] for _ = 1:nthreads]
+        saxs = SplitAxes(axsimg, nthreads - 0.2)   # give main thread less work since it also schedules the others
+        for i = nthreads:-1:2
+            push!(tasks, Threads.@spawn computeft!(F, img, saxs[i], CartesianIndex(()), weights, tmps[i]))
+        end
+        computeft!(F, img, saxs[1], CartesianIndex(()), weights, tmps[1])
+        while !isempty(tasks)
+            wait(pop!(tasks))
+        end
+        # Finish the last dimension
+        saxs1 = SplitAxes(axsimg[1:N-1], nthreads - 0.2)
+        for i = nthreads:-1:2
+            push!(tasks, Threads.@spawn finishft!(F, img, (saxs1[i]..., axsimg[end]), CartesianIndex(()), weights, tmps[i]))
+        end
+        finishft!(F, img, (saxs1[1]..., axsimg[end]), CartesianIndex(()), weights, tmps[1])
+        while !isempty(tasks)
+            wait(pop!(tasks))
+        end
+    end
+    return F
 end
 
 """
@@ -68,33 +97,38 @@ function distance_transform(F::AbstractArray{CartesianIndex{N},N}, w::Union{Noth
     D
 end
 
-function computeft!(F, I, jpost::CartesianIndex{K}, pixelspacing, tmp) where K
+function computeft!(F, img, axsimg, jpost::CartesianIndex{K}, pixelspacing, tmp) where K
     # tmp is workspace for voronoift!
     ∅ = nullindex(F)                             # sentinel position
-    if K == ndims(I)-1                           # innermost loop (d=1 case, line 1)
+    if K == ndims(img)-1                           # innermost loop (d=1 case, line 1)
         # Fig. 2, lines 2-8
-        @inbounds @simd for i1 in axes(I, 1)
-            F[i1, jpost] = I[i1, jpost] ? CartesianIndex(i1, jpost) : ∅
+        @inbounds @simd for i1 in axes(img, 1)
+            F[i1, jpost] = gray(img[i1, jpost]) ? CartesianIndex(i1, jpost) : ∅
         end
     else                                         # recursively handle trailing dimensions
         # Fig. 2, lines 10-12
-        for i1 in axes(I, ndims(I) - K)
-            computeft!(F, I, CartesianIndex(i1, jpost), pixelspacing, tmp)
+        for i1 in axsimg[ndims(img) - K]
+            computeft!(F, img, axsimg, CartesianIndex(i1, jpost), pixelspacing, tmp)
         end
     end
+    K == 0 && return F                           # for the final axis, threads will be split across next-to-last axis
+    return finishft!(F, img, axsimg, jpost, pixelspacing, tmp)
+end
+
+function finishft!(F, img, axsimg, jpost, pixelspacing, tmp)
     # Fig. 2, lines 14-20
-    axespre = truncatet(axes(F), jpost)          # discard the trailing axes of F
+    axespre = truncatet(axsimg, jpost)          # first N-K-1 axes (these are "finished" within each K+1-dimensional slice)
     for jpre in CartesianIndices(axespre)
-        voronoift!(F, I, jpre, jpost, pixelspacing, tmp)
+        voronoift!(F, img, jpre, jpost, pixelspacing, tmp)     # finish axis N-K in K-dimensional slice `jpost`
     end
-    F
+    return F
 end
 
-function voronoift!(F, I, jpre, jpost, pixelspacing, tmp)
-    d = length(jpre)+1
+function voronoift!(F, img, jpre, jpost, pixelspacing, tmp)
+    d = length(jpre)+1   # axis to work along
     ∅ = nullindex(F)
     empty!(tmp)
-    for i in axes(I, d)
+    for i in axes(img, d)
         # Fig 3, lines 3-13
         xi = CartesianIndex(jpre, i, jpost)
         @inbounds fi = F[xi]
@@ -115,7 +149,7 @@ function voronoift!(F, I, jpre, jpost, pixelspacing, tmp)
     # Fig 3, lines 18-24
     l = 1
     @inbounds fthis = tmp[l].fi
-    for i in axes(I, d)
+    for i in axes(img, d)
         xi = CartesianIndex(jpre, i, jpost)
         d2this = wnorm2(xi-fthis, pixelspacing)
         while l < nS

src/deprecations.jl

@@ -4,3 +4,6 @@
 @deprecate label_components!(out::AbstractArray{Int,N}, A::AbstractArray{T,N}, connectivity::Array{Bool,N}, bkg) where {T,N}  label_components!(out, A, connectivity; bkg=bkg)
 
 @deprecate imfill(img::AbstractArray{Bool}, interval::Tuple{Real,Real}, dims::Union{Dims, AbstractVector{Int}}) imfill(img, interval; dims=dims)
+
+import .FeatureTransform: feature_transform
+@deprecate feature_transform(img, weights; kwargs...) feature_transform(img; weights=weights, kwargs...)

test/bwdist.jl

@@ -68,6 +68,8 @@
         @test F == ind2cart([1 1 1; 2 2 13; 2 8 13; 4 9 14; 4 10 10])
         D = distance_transform(F)
         @test D == [0.0 1.0 2.0; 0.0 1.0 1.0; 1.0 0.0 0.0; 0.0 0.0 0.0; 1.0 0.0 1.0]
+
+        @test feature_transform(Gray.(A)) == F
     end
 
     @testset "Corner Case Images" begin