JuliaImages / ImageSegmentation.jl
1
import Base.isless
2

3
struct PixelKey{CT, N}
4 12
    val::CT
5
    time_step::Int
6
    source::CartesianIndex{N}
7
end
8 12
isless(a::PixelKey, b::PixelKey) = (a.val < b.val) || (a.val == b.val && a.time_step < b.time_step)
9

10
"""Calculate the euclidean distance between two `CartesianIndex` structs"""
11 12
@inline _euclidean(a::CartesianIndex{N}, b::CartesianIndex{N}) where {N} = sqrt(sum(Tuple(a - b) .^ 2))
12

13
"""
14
```
15
segments                = watershed(img, markers; compactness, mask)
16
```
17
Segments the image using watershed transform. Each basin formed by watershed transform corresponds to a segment.
18
If you are using image local minimas as markers, consider using [`hmin_transform`](@ref) to avoid oversegmentation.
19

20
Parameters:
21
-    img            = input grayscale image
22
-    markers        = An array (same size as img) with each region's marker assigned a index starting from 1. Zero means not a marker.
23
                      If two markers have the same index, their regions will be merged into a single region.
24
                      If you have markers as a boolean array, use `label_components`.
25
-    compactness    = Use the compact watershed algorithm with the given compactness parameter. Larger values lead to more regularly
26
                      shaped watershed basins.[^1]
27
-    mask           = Only segment pixels where the value of `mask` is true, used to restrict segmentation to only areas of interest
28

29

30
[^1]: https://www.tu-chemnitz.de/etit/proaut/publications/cws_pSLIC_ICPR.pdf
31

32
# Example
33

34
```jldoctest; setup = :(using Images, ImageSegmentation)
35
julia> seeds = falses(100, 100); seeds[50, 25] = true; seeds[50, 75] = true;
36

37
julia> dists = distance_transform(feature_transform(seeds)); # calculate distances from seeds
38

39
julia> markers = label_components(seeds); # give each seed a unique integer id
40

41
julia> results = watershed(dists, markers);
42

43
julia> labels_map(results); # labels of segmented image
44
```
45
"""
46 6
function watershed(img::AbstractArray{T, N},
47
                   markers::AbstractArray{S,N};
48
                   mask::AbstractArray{Bool, N}=fill(true, axes(img)),
49
                   compactness::Real = 0.0) where {T<:Images.NumberLike, S<:Integer, N}
50

51 12
    if axes(img) != axes(markers)
52 0
        error("image size doesn't match marker image size")
53 12
    elseif axes(img) != axes(mask)
54 0
        error("image size doesn't match mask size")
55
    end
56

57 6
    compact = compactness > 0.0
58 6
    segments = copy(markers)
59 12
    PK = PixelKey{compact ? floattype(T) : T, N}
60 12
    pq = PriorityQueue{CartesianIndex{N}, PK}()
61 6
    time_step = 0
62

63 6
    R = CartesianIndices(axes(img))
64 6
    Istart, Iend = first(R), last(R)
65 12
    for i in R
66 12
        if markers[i] > 0
67 12
            for j in CartesianIndices(_colon(max(Istart,i-_oneunit(i)), min(i+_oneunit(i),Iend)))
68 12
                if segments[j] == 0
69 6
                    segments[j] = markers[i]
70 6
                    enqueue!(pq, j, PK(img[i], time_step, j))
71 12
                    time_step += 1
72
                end
73
            end
74
        end
75
    end
76

77 12
    while !isempty(pq)
78 12
        curr_idx, curr_elem = dequeue_pair!(pq)
79 6
        segments_current = segments[curr_idx]
80

81
        # If we're using the compact algorithm, we need assign grouping for a given location
82
        # when it comes off the queue since we could find a better suited watershed later.
83 12
        if compact
84 12
            if segments_current > 0 && curr_idx != curr_elem.source
85
                # this is a non-marker location that we've already assigned
86 0
                continue
87
            end
88
            # group this location with its watershed
89 6
            segments[curr_idx] = segments[curr_elem.source]
90
        end
91

92 6
        img_current = img[curr_idx]
93 12
        for j in CartesianIndices(_colon(max(Istart,curr_idx-_oneunit(curr_idx)), min(curr_idx+_oneunit(curr_idx),Iend)))
94

95
            # if this location is false in the mask, we skip it
96 12
            (!mask[j]) && continue
97
            # only continue if this is a position that we haven't assigned yet
98 12
            if segments[j] == 0
99
                # if we're doing a simple watershed, we can go ahead and set the final grouping for a new
100
                # ungrouped position the moment we first encounter it
101 12
                if !compact
102 6
                    segments[j] = segments_current
103 12
                    new_value = img_current
104
                else
105
                    # in the compact algorithm case, we don't set the grouping
106
                    # at push-time and instead calculate a temporary value based
107
                    # on the weighted sum of the intensity and distance to the
108
                    # current source marker.
109 6
                    new_value = floattype(T)(img_current + compactness * _euclidean(j, curr_elem.source))
110
                end
111

112
                # if this position is in the queue and we're using the compact algorithm, we need to replace
113
                # its watershed if we find one that it better belongs to
114 12
                if compact && j in keys(pq)
115 6
                    elem = pq[j]
116 6
                    new_elem = PK(new_value, time_step, curr_elem.source)
117

118 12
                    if new_elem < elem
119 12
                        pq[j] = new_elem # update the watershed
120 12
                        time_step += 1
121
                    end
122
                else
123

124 12
                    pq[j] = PK(new_value, time_step, curr_elem.source)
125 12
                    time_step += 1
126
                end
127
            end
128
        end
129
    end
130

131 6
    TM = meantype(T)
132 6
    num_segments        = Int(maximum(segments))
133 6
    labels              = Array(1:num_segments)
134 12
    region_means        = Dict{Int, TM}()
135 12
    region_pix_count    = Dict{Int, Int}()
136

137 12
    for i in R
138 12
        region_pix_count[segments[i]] = get(region_pix_count, segments[i], 0) + 1
139 12
        region_means[segments[i]] = get(region_means, segments[i], zero(TM)) + (img[i] - get(region_means, segments[i], zero(TM)))/region_pix_count[segments[i]]
140
    end
141 12
    return SegmentedImage(segments, labels, region_means, region_pix_count)
142
end
143

144
"""
145
```
146
out = hmin_transform(img, h)
147
```
148
Suppresses all minima in grayscale image whose depth is less than h.
149

150
H-minima transform is defined as the reconstruction by erosion of (img + h) by img. See Morphological image analysis by Soille pg 170-172.
151
"""
152 6
function hmin_transform(img::Array{T, N}, h::Real) where {T<:Images.NumberLike, N}
153 12
    out = img.+h
154 12
    while true
155 6
        temp = max.(img, erode(out))
156 12
        if temp == out
157 12
            break
158
        end
159 12
        out = temp
160
    end
161 12
    return out
162
end

Read our documentation on viewing source code .

Loading