Add support for compact and masked watersheds (#33)

Author: tlnagy

Project.toml

@@ -1,6 +1,6 @@
 name = "ImageSegmentation"
 uuid = "80713f31-8817-5129-9cf8-209ff8fb23e1"
-version = "1.4.2"
+version = "1.4.3"
 
 [deps]
 Clustering = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5"

src/watershed.jl

@@ -1,14 +1,18 @@
 import Base.isless
 
-struct PixelKey{CT}
+struct PixelKey{CT, N}
     val::CT
     time_step::Int
+    source::CartesianIndex{N}
 end
-isless(a::PixelKey{T}, b::PixelKey{T}) where {T} = (a.val < b.val) || (a.val == b.val && a.time_step < b.time_step)
+isless(a::PixelKey, b::PixelKey) = (a.val < b.val) || (a.val == b.val && a.time_step < b.time_step)
+
+"""Calculate the euclidean distance between two `CartesianIndex` structs"""
+@inline _euclidean(a::CartesianIndex{N}, b::CartesianIndex{N}) where {N} = sqrt(sum(Tuple(a - b) .^ 2))
 
 """
 ```
-segments                = watershed(img, markers)
+segments                = watershed(img, markers; compactness, mask)
 ```
 Segments the image using watershed transform. Each basin formed by watershed transform corresponds to a segment.
 If you are using image local minimas as markers, consider using [`hmin_transform`](@ref) to avoid oversegmentation.
@@ -18,18 +22,42 @@ Parameters:
 -    markers        = An array (same size as img) with each region's marker assigned a index starting from 1. Zero means not a marker.
                       If two markers have the same index, their regions will be merged into a single region.
                       If you have markers as a boolean array, use `label_components`.
+-    compactness    = Use the compact watershed algorithm with the given compactness parameter. Larger values lead to more regularly
+                      shaped watershed basins.[^1]
+-    mask           = Only segment pixels where the value of `mask` is true, used to restrict segmentation to only areas of interest
+
+
+[^1]: https://www.tu-chemnitz.de/etit/proaut/publications/cws_pSLIC_ICPR.pdf
+
+# Example
+
+```jldoctest; setup = :(using Images, ImageSegmentation)
+julia> seeds = falses(100, 100); seeds[50, 25] = true; seeds[50, 75] = true;
 
+julia> dists = distance_transform(feature_transform(seeds)); # calculate distances from seeds
 
+julia> markers = label_components(seeds); # give each seed a unique integer id
 
+julia> results = watershed(dists, markers);
+
+julia> labels_map(result); # labels of segmented image
+```
 """
-function watershed(img::AbstractArray{T, N}, markers::AbstractArray{S,N}) where {T<:Images.NumberLike, S<:Integer, N}
+function watershed(img::AbstractArray{T, N},
+                   markers::AbstractArray{S,N};
+                   mask::AbstractArray{Bool, N}=fill(true, axes(img)),
+                   compactness::Real = 0.0) where {T<:Images.NumberLike, S<:Integer, N}
 
     if axes(img) != axes(markers)
         error("image size doesn't match marker image size")
+    elseif axes(img) != axes(mask)
+        error("image size doesn't match mask size")
     end
 
+    compact = compactness > 0.0
     segments = copy(markers)
-    pq = PriorityQueue{CartesianIndex{N}, PixelKey{T}}()
+    PK = PixelKey{compact ? floattype(T) : T, N}
+    pq = PriorityQueue{CartesianIndex{N}, PK}()
     time_step = 0
 
     R = CartesianIndices(axes(img))
@@ -39,22 +67,63 @@ function watershed(img::AbstractArray{T, N}, markers::AbstractArray{S,N}) where
             for j in CartesianIndices(_colon(max(Istart,i-one(i)), min(i+one(i),Iend)))
                 if segments[j] == 0
                     segments[j] = markers[i]
-                    enqueue!(pq, j, PixelKey(img[i], time_step))
+                    enqueue!(pq, j, PK(img[i], time_step, j))
                     time_step += 1
                 end
             end
         end
     end
 
     while !isempty(pq)
-        current = dequeue!(pq)
-        segments_current = segments[current]
-        img_current = img[current]
-        for j in CartesianIndices(_colon(max(Istart,current-one(current)), min(current+one(current),Iend)))
+        curr_idx, curr_elem = dequeue_pair!(pq)
+        segments_current = segments[curr_idx]
+
+        # If we're using the compact algorithm, we need assign grouping for a given location
+        # when it comes off the queue since we could find a better suited watershed later.
+        if compact
+            if segments_current > 0 && curr_idx != curr_elem.source
+                # this is a non-marker location that we've already assigned
+                continue
+            end
+            # group this location with its watershed
+            segments[curr_idx] = segments[curr_elem.source]
+        end
+
+        img_current = img[curr_idx]
+        for j in CartesianIndices(_colon(max(Istart,curr_idx-one(curr_idx)), min(curr_idx+one(curr_idx),Iend)))
+
+            # if this location is false in the mask, we skip it
+            (!mask[j]) && continue
+            # only continue if this is a position that we haven't assigned yet
             if segments[j] == 0
-                segments[j] = segments_current
-                enqueue!(pq, j, PixelKey(img_current, time_step))
-                time_step += 1
+                # if we're doing a simple watershed, we can go ahead and set the final grouping for a new
+                # ungrouped position the moment we first encounter it
+                if !compact
+                    segments[j] = segments_current
+                    new_value = img_current
+                else
+                    # in the compact algorithm case, we don't set the grouping
+                    # at push-time and instead calculate a temporary value based
+                    # on the weighted sum of the intensity and distance to the
+                    # current source marker.
+                    new_value = floattype(T)(img_current + compactness * _euclidean(j, curr_elem.source))
+                end
+
+                # if this position is in the queue and we're using the compact algorithm, we need to replace
+                # its watershed if we find one that it better belongs to
+                if compact && j in keys(pq)
+                    elem = pq[j]
+                    new_elem = PK(new_value, time_step, curr_elem.source)
+
+                    if new_elem < elem
+                        pq[j] = new_elem # update the watershed
+                        time_step += 1
+                    end
+                else
+
+                    pq[j] = PK(new_value, time_step, curr_elem.source)
+                    time_step += 1
+                end
             end
         end
     end

test/watershed.jl

@@ -11,24 +11,56 @@ using ImageFiltering
     markers[1, 1] = 1
     markers[50, 50] = 2
 
-    result = watershed(img, markers)
+    @testset "classic watershed" begin
+        # with classic watershed, we expect the top left label should spread around
+        # the center label.
+        result = watershed(img, markers)
 
-    @test length(result.segment_labels) == 2
-    @test result.segment_labels == collect(1:2)
-    @test all(label->(label==result.image_indexmap[50,50]), result.image_indexmap[26:74,26:74])
+        @test length(result.segment_labels) == 2
+        @test result.segment_labels == collect(1:2)
+        @test all(label->(label==result.image_indexmap[50,50]), result.image_indexmap[26:74,26:74])
+    end
 
+    @testset "masked watershed" begin
+        mask = trues(size(img))
+        mask[60:70, 60:70] .= false
 
-    img = ones(15, 15)
-    #minima of depth 0.2
-    img[3:5, 3:5] .= 0.9
-    img[4,4] = 0.8
-    #minima of depth 0.7
-    img[9:13, 9:13] .= 0.8
-    img[10:12, 10:12] .= 0.7
-    img[11,11] = 0.3
+        result = watershed(img, markers, mask=mask)
+        labels = labels_map(result)
 
-    out = hmin_transform(img, 0.25)
+        # where the mask is false, no label should be assigned
+        @test sum(labels[.~ mask]) == 0
 
-    @test findlocalminima(img) == [CartesianIndex(4, 4), CartesianIndex(11, 11)]
-    @test findlocalminima(out) == [CartesianIndex(11, 11)]
+        result = watershed(img, markers, compactness=10.0, mask=mask)
+        labels = labels_map(result)
+
+        # where the mask is false, no label should be assigned
+        @test sum(labels[.~ mask]) == 0
+    end
+
+    @testset "compact watershed" begin
+        result = watershed(img, markers, compactness=10.0)
+        labels = labels_map(result)
+
+        # since this is using the compact algorithm with a high value for
+        # compactness, the boundary between labels 1 and 2 should occur halfway
+        # between the two markers
+        @test sum(labels .== 1) == sum(1:50) - 2
+    end
+
+    @testset "h-minima transform" begin
+        img = ones(15, 15)
+        #minima of depth 0.2
+        img[3:5, 3:5] .= 0.9
+        img[4,4] = 0.8
+        #minima of depth 0.7
+        img[9:13, 9:13] .= 0.8
+        img[10:12, 10:12] .= 0.7
+        img[11,11] = 0.3
+
+        out = hmin_transform(img, 0.25)
+
+        @test findlocalminima(img) == [CartesianIndex(4, 4), CartesianIndex(11, 11)]
+        @test findlocalminima(out) == [CartesianIndex(11, 11)]
+    end
 end