JuliaImages / ImageSegmentation.jl
1
"""
2
```
3
segments                = meanshift(img, spatial_radius, range_radius; iter=50, eps=0.01))
4
```
5
Segments the image using meanshift clustering. Returns a `SegmentedImage`.
6

7
Parameters:
8
-    img                            = input grayscale image
9
-    spatial_radius/range_radius    = bandwidth parameters of truncated normal kernel. Controlling the size of the kernel determines the resolution of the mode detection.
10
-    iter/eps                       = stopping criterion for meanshift procedure. The algorithm stops after iter iterations or if kernel center moves less than eps distance in an update step, whichever comes first.
11
"""
12 3
function meanshift(img::Array{CT, 2}, spatial_radius::Real, range_radius::Real; iter::Int = 50, eps::Real = 0.01) where CT
13

14 6
    rows, cols = size(img)
15 3
    rowbins = Int(floor(rows/spatial_radius))
16 3
    colbins = Int(floor(cols/spatial_radius))
17 3
    colorbins = Int(floor(1/range_radius))
18

19 6
    colorbin(x)::Int = min(floor(x/range_radius) + 1, colorbins)
20 6
    rowbin(x)::Int = min(floor(x/spatial_radius) +1, rowbins)
21 6
    colbin(x)::Int = min(floor(x/spatial_radius) +1, colbins)
22

23 3
    buckets = Array{Vector{CartesianIndex{2}}}(undef, rowbins, colbins, colorbins)
24 6
    for i in CartesianIndices(size(buckets))
25 6
        buckets[i]=Array{CartesianIndex{2}}(undef, 0)
26
    end
27

28 6
    for i in CartesianIndices(size(img))
29 6
        push!( buckets[rowbin(i[1]), colbin(i[2]), colorbin(img[i])], i)
30
    end
31

32 6
    function dist2(a::SVector, b::SVector)::Float64
33 6
        return ((a[1]-b[1])/spatial_radius)^2 + ((a[2]-b[2])/spatial_radius)^2 + ((a[3]-b[3])/range_radius)^2
34
    end
35

36 6
    function neighborhood_mean(pt::SVector{3, Float64})::SVector{3, Float64}
37 6
        den = 0.0
38 3
        num = SVector(0.0, 0.0, 0.0)
39

40 6
        R = CartesianIndices(size(buckets))
41 3
        I1, Iend = first(R), last(R)
42 6
        I = CartesianIndex(rowbin(pt[1]), colbin(pt[2]), colorbin(pt[3]))
43

44 6
        for J in _colon(max(I1, I-_oneunit(I)), min(Iend, I+_oneunit(I)))
45 6
            for point in buckets[J]
46 6
                if dist2(pt, SVector(point[1], point[2], img[point])) <= 1
47 6
                    num += SVector(point[1], point[2], img[point])
48 6
                    den += 1
49
                end
50
            end
51
        end
52

53 6
        return den<=0 ? pt : num/den
54
    end
55

56 3
    modes = Array{Float64}(undef, 3, rows*cols)
57

58 6
    for i in CartesianIndices(size(img))
59 3
        pt::SVector{3, Float64} = SVector(i[1], i[2], img[i])
60 6
        for j in 1:iter
61 6
            nextpt = neighborhood_mean(pt)
62 6
            if dist2(pt, nextpt) < eps^2
63 6
                break
64
            end
65 6
            pt = nextpt
66
        end
67 6
        modes[:, i[1] + (i[2]-1)*rows] = [pt[1]/spatial_radius, pt[2]/spatial_radius, pt[3]/range_radius]
68
    end
69

70 6
    clusters = dbscan(modes, 1.414)
71 3
    num_segments = length(clusters)
72 3
    TM = meantype(CT)
73 3
    result              = similar(img, Int)
74 3
    labels              = Array(1:num_segments)
75 6
    region_means        = Dict{Int, TM}()
76 6
    region_pix_count    = Dict{Int, Int}()
77

78 3
    cluster_idx = 0
79 6
    for cluster in clusters
80 3
        cluster_idx += 1
81 6
        region_pix_count[cluster_idx] = cluster.size
82 6
        for index in [cluster.core_indices; cluster.boundary_indices]
83 3
            i, j = (index-1)%rows + 1, floor(Int, (index-1)/rows) + 1
84 3
            result[i, j] = cluster_idx
85 6
            region_means[cluster_idx] = get(region_means, cluster_idx, zero(TM)) + img[i, j]
86
        end
87 6
        region_means[cluster_idx] /= region_pix_count[cluster_idx]
88
    end
89

90 6
    return SegmentedImage(result, labels, region_means, region_pix_count)
91
end

Read our documentation on viewing source code .

Loading