Skip to content

Commit 760a4fe

Browse files
authored
graph-based segmentation merge (#64)
1 parent 15ea1ff commit 760a4fe

File tree

5 files changed

+330
-4
lines changed

5 files changed

+330
-4
lines changed

Project.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
1010
Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
1111
LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
1212
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
13+
MetaGraphs = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
1314
RegionTrees = "dee08c22-ab7f-5625-9660-a9af2021b33f"
1415
SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
1516
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -23,6 +24,7 @@ Documenter = "0.24, 0.25"
2324
ImageFiltering = "0.6"
2425
Images = "0.18, 0.19, 0.20, 0.21, 0.22, 0.23"
2526
LightGraphs = "1.1"
27+
MetaGraphs = "0.6.6"
2628
RegionTrees = "0.2, 0.3"
2729
SimpleWeightedGraphs = "1"
2830
StaticArrays = "0.9, 0.10, 0.11, 0.12, 1.0"

src/ImageSegmentation.jl

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module ImageSegmentation
33
import Base: show
44

55
using LinearAlgebra, Statistics
6-
using Images, DataStructures, StaticArrays, ImageFiltering, LightGraphs, SimpleWeightedGraphs, RegionTrees, Distances, StaticArrays, Clustering
6+
using Images, DataStructures, StaticArrays, ImageFiltering, LightGraphs, SimpleWeightedGraphs, RegionTrees, Distances, StaticArrays, Clustering, MetaGraphs
77
import Clustering: kmeans, fuzzy_cmeans
88

99
include("compat.jl")
@@ -15,6 +15,7 @@ include("watershed.jl")
1515
include("region_merging.jl")
1616
include("meanshift.jl")
1717
include("clustering.jl")
18+
include("merge_segments.jl")
1819

1920
export
2021
#accessor methods
@@ -39,7 +40,8 @@ export
3940
meanshift,
4041
kmeans,
4142
fuzzy_cmeans,
42-
43+
merge_segments,
44+
4345
# types
4446
SegmentedImage,
4547
ImageEdge

src/merge_segments.jl

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
"""
2+
seg2 = merge_segments(seg, threshold)
3+
4+
Merges segments in a [`SegmentedImage`](@ref) by building a region adjacency
5+
graph (RAG) and merging segments connected by edges with weight less than
6+
`threshold`.
7+
8+
# Arguments:
9+
* `seg` : SegmentedImage to be merged.
10+
* `threshold` : Upper bound of the adjacent segment color difference to
11+
consider merging segments.
12+
13+
# Citation:
14+
Vighnesh Birodkar
15+
"Hierarchical merging of region adjacency graphs"
16+
https://vcansimplify.wordpress.com/2014/08/17/hierarchical-merging-of-region-adjacency-graphs/
17+
"""
18+
function merge_segments(seg::SegmentedImage, threshold::Number)::SegmentedImage
19+
g = seg_to_graph(seg)
20+
21+
# Populate a heap of all the edges, and a Bool indicating whether the edge
22+
# is valid. All edges are initially valid. The reason for this is that heap
23+
# removal would be expensive, so instead, we invalidate the edge entry in the
24+
# heap.
25+
function weight(t::Tuple{Edge{Int}, Bool})::Real
26+
return has_prop(g, t[1], :weight) ? get_prop(g, t[1], :weight) : 0
27+
end
28+
29+
edge_heap = MutableBinaryHeap{Tuple{Edge{Int}, Bool}}(Base.By(weight),
30+
[(e, true) for e in edges(g)]
31+
)
32+
sizehint!(edge_heap, 3 * length(edge_heap)) # Overkill, or not enough?
33+
for n in edge_heap.nodes
34+
set_prop!(g, n.value[1], :handle, n.handle)
35+
end
36+
37+
# Merge all edges less than threshold
38+
while !isempty(edge_heap) && weight(first(edge_heap)) < threshold
39+
e, valid = pop!(edge_heap)
40+
if valid
41+
# Invalidate all edges touching this edge.
42+
invalidate_neighbors!(edge_heap, g, e)
43+
44+
# Merge the two nodes into one (keep e.dst, obsolete e.src)
45+
merge_node_props!(g, e)
46+
47+
# Make new edges to the merged node.
48+
new_edges = add_neighboring_edges!(g, e)
49+
50+
# Remove edges to src.
51+
# Don't call rem_vertex!(g, e.src); it would renumber all vertices.
52+
for n in collect(neighbors(g, e.src))
53+
rem_edge!(g, e.src, n)
54+
end
55+
56+
# Add new edges to heap.
57+
for e in new_edges
58+
handle = push!(edge_heap, (e, true))
59+
set_prop!(g, e, :handle, handle)
60+
end
61+
end
62+
end
63+
64+
return resegment(seg, g)
65+
end
66+
67+
68+
"""
69+
g = seg_to_graph(seg)
70+
71+
Given a [`SegmentedImage`](@ref), produces a region adjacency [`MetaGraph`](@ref)
72+
and stores segment metadata on the vertices. Edge weight is determined by
73+
color difference.
74+
75+
# Arguments:
76+
* `seg` : a [`SegmentedImage`](@ref)
77+
"""
78+
function seg_to_graph(seg::SegmentedImage)::MetaGraph
79+
weight(i, j) = colordiff(segment_mean(seg, i), segment_mean(seg, j))
80+
rag, _ = region_adjacency_graph(seg, weight)
81+
82+
g = MetaGraph(rag)
83+
for v in vertices(rag)
84+
set_prop!(g, v, :labels, [v])
85+
set_prop!(g, v, :pixel_count, seg.segment_pixel_count[v])
86+
set_prop!(g, v, :mean_color, seg.segment_means[v])
87+
set_prop!(g, v, :total_color, seg.segment_means[v] * seg.segment_pixel_count[v])
88+
end
89+
90+
for e in edges(rag)
91+
set_prop!(g, Edge(e.src, e.dst), :weight, e.weight)
92+
end
93+
return g
94+
end
95+
96+
97+
"""
98+
seg2 = resegment(seg1, rag)
99+
100+
Takes a segmentation and a region adjacency graph produced by `merge_segments`
101+
and produces an segmentation that corresponds to the graph.
102+
103+
# Arguments:
104+
* `seg` : a [`SegmentedImage`](@ref)
105+
* `g` : a [`MetaGraph`](@ref) representing a merged Region Adjacency
106+
Graph
107+
"""
108+
function resegment(seg::SegmentedImage, g::MetaGraph)::SegmentedImage
109+
# Find all the vertices of g that remain post-merge.
110+
remaining = collect(filter(v -> 0 < length(props(g, v)), vertices(g)))
111+
112+
px_labels = copy(seg.image_indexmap)
113+
# Re-label all pixels with the vertex they were merged to.
114+
for v in remaining
115+
labels = get_prop(g, v, :labels)
116+
for l in labels
117+
if l != v
118+
ix = findall(x -> x == l, px_labels)
119+
px_labels[ix] .= v
120+
end
121+
end
122+
end
123+
124+
# Re-number our labels so that they are dense (no gaps) and
125+
# construct the other objects SegmentedImage needs.
126+
means, px_counts = Dict{Int, Colorant}(), Dict{Int, Int}()
127+
for (i, v) in enumerate(remaining)
128+
ix = findall(x -> x == v, px_labels)
129+
px_labels[ix] .= i
130+
means[i] = get_prop(g, v, :mean_color)
131+
px_counts[i] = get_prop(g, v, :pixel_count)
132+
end
133+
134+
labels = collect(1:length(remaining))
135+
136+
return SegmentedImage(px_labels, labels, means, px_counts)
137+
end
138+
139+
140+
"""
141+
merge_node_props!(g, e)
142+
143+
Takes edge `e` in [`MetaGraph`](@ref) `g` and merges the props from its `src`
144+
and `dst` into its `dst`, clearing all props from `src`.
145+
146+
"""
147+
function merge_node_props!(g::MetaGraph, e::AbstractEdge)
148+
src, dst = e.src, e.dst
149+
clr = get_prop(g, dst, :total_color) + get_prop(g, src, :total_color)
150+
npx = get_prop(g, dst, :pixel_count) + get_prop(g, src, :pixel_count)
151+
152+
set_prop!(g, dst, :total_color, clr)
153+
set_prop!(g, dst, :pixel_count, npx)
154+
set_prop!(g, dst, :mean_color, clr / npx)
155+
set_prop!(g, dst, :labels, vcat(
156+
get_prop(g, src, :labels),
157+
get_prop(g, dst, :labels)
158+
))
159+
160+
# Clear props on the now unused node src, to make its obsolescence clear.
161+
clear_props!(g, src)
162+
end
163+
164+
165+
"""
166+
new_edges = add_neighboring_edges!(g, e)
167+
168+
Finds the nodes neighboring `e` in graph `g`, creates edges from them to its
169+
`dst`, and sets the weight of the new edges.
170+
171+
# Arguments:
172+
* `g` : a [`MetaGraph`](@ref)
173+
* `e` : an [`AbstractEdge`](@ref)
174+
"""
175+
function add_neighboring_edges!(g::MetaGraph, e::AbstractEdge)
176+
edges = Edge{eltype(e)}[]
177+
edge_neighbors = union(Set(neighbors(g, e.src)), Set(neighbors(g, e.dst)))
178+
for n in setdiff(edge_neighbors, e.src, e.dst)
179+
edge = Edge(e.dst, n)
180+
add_edge!(g, edge)
181+
set_prop!(g, edge, :weight, _weight_mean_color(g, edge))
182+
push!(edges, edge)
183+
end
184+
185+
return edges
186+
end
187+
188+
189+
"""
190+
invalidate_neighbors!(edge_heap, g, e)
191+
192+
Finds the neighbors of `e` in graph `g` and invalidates them in `edge_heap`.
193+
194+
# Arguments:
195+
* `edge_heap` : a [`MutableBinaryHeap`](@ref)
196+
* `g` : a [`MetaGraph`](@ref)
197+
* `e` : an [`AbstractEdge`](@ref)
198+
"""
199+
function invalidate_neighbors!(edge_heap::MutableBinaryHeap, g::MetaGraph, e::AbstractEdge)
200+
function invalidate(src, dst)
201+
for n in setdiff(Set(neighbors(g, src)), dst)
202+
edge = Edge(src, n)
203+
h = get_prop(g, edge, :handle)
204+
update!(edge_heap, h, (edge, false))
205+
end
206+
end
207+
invalidate(e.src, e.dst)
208+
invalidate(e.dst, e.src)
209+
end
210+
211+
212+
"""
213+
weight = _weight_mean_color(g, v1, v2))
214+
215+
Compute the weight of an edge in [`MetaGraph`](@ref) `g` as the difference
216+
in mean colors of each vertex.
217+
218+
# Arguments:
219+
* `g` : a [`MetaGraph`](@ref)
220+
* `e` : an [`AbstractEdge`](@ref)
221+
"""
222+
function _weight_mean_color(g::MetaGraph, e::AbstractEdge)::Real
223+
return colordiff(
224+
get_prop(g, e.src, :mean_color),
225+
get_prop(g, e.dst, :mean_color)
226+
)
227+
end
228+

test/merge_segments.jl

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
include("../src/merge_segments.jl")
2+
3+
@testset "SegmentedImage Merge" begin
4+
# Set up 3 variables: test_graph, img, seg for use in tests
5+
g = MetaGraph(6)
6+
7+
for v in 1:6
8+
set_prop!(g, v, :total_color, v)
9+
set_prop!(g, v, :pixel_count, v)
10+
set_prop!(g, v, :mean_color, Gray(Float64(v)))
11+
set_prop!(g, v, :labels, v)
12+
end
13+
14+
for i in 1:5
15+
add_edge!(g, i, i+1, :weight, i + i + 1)
16+
end
17+
18+
img = fill(1.0, (10, 10))
19+
img[:, 5:10] .= 2.0
20+
img = map(Gray, img)
21+
seg = fast_scanning(img, 0.5)
22+
test_graph = copy(g)
23+
# end setup
24+
25+
# Test 1. Merge 2 nodes
26+
g = copy(test_graph)
27+
merge_node_props!(g, Edge(1, 2))
28+
29+
# Check that labels got properly merged
30+
@test sort(get_prop(g, 2, :labels)) == [1, 2]
31+
32+
# Check that the other node props were properly updated
33+
@test get_prop(g, 2, :total_color) == 3
34+
@test get_prop(g, 2, :pixel_count) == 3
35+
@test get_prop(g, 2, :mean_color) == 1
36+
37+
# Ensure props were cleared
38+
@test length(props(g, 1)) == 0
39+
40+
# Ensure no weights were affected.
41+
for v in 1:5
42+
@test get_prop(g, Edge(v, v+1), :weight) == (v + v + 1)
43+
end
44+
45+
# Test 2. Merge all nodes
46+
g = copy(test_graph)
47+
for i in 1:5
48+
merge_node_props!(g, Edge(i, i+1))
49+
end
50+
51+
@test sort(get_prop(g, 6, :labels)) == 1:6
52+
@test get_prop(g, 6, :total_color) == sum(1:6)
53+
@test get_prop(g, 6, :pixel_count) == sum(1:6)
54+
@test get_prop(g, 6, :mean_color) == 1
55+
56+
for i in 1:5
57+
@test length(props(g, i)) == 0
58+
end
59+
60+
# Test 3. add_neighboring_edges!
61+
g = copy(test_graph)
62+
add_edge!(g, 1, 6, :weight, 7) # setup
63+
64+
added = add_neighboring_edges!(g, Edge(1, 2))
65+
66+
@test length(added) == 2
67+
@test has_edge(g, Edge(2, 3))
68+
@test has_edge(g, Edge(2, 6))
69+
70+
# Test 4. seg to graph
71+
g = seg_to_graph(seg)
72+
@test length(edges(g)) == 1
73+
@test length(vertices(g)) == 2
74+
@test has_edge(g, Edge(1, 2))
75+
@test haskey(props(g, Edge(1, 2)), :weight)
76+
@test get_prop(g, Edge(1, 2), :weight) Colors.colordiff(1.0, 2.0)
77+
78+
79+
# Test 5. resegment
80+
g = seg_to_graph(seg)
81+
82+
# identity resegment
83+
seg2 = resegment(seg, g)
84+
for f in [labels_map, segment_labels, segment_pixel_count, segment_mean]
85+
@test f(seg2) == f(seg)
86+
end
87+
88+
89+
# Test 6. merge
90+
seg2 = merge_segments(seg, 40) # 40 > colordiff(1.0, 2.0)
91+
@test segment_labels(seg2) == [1]
92+
end

test/runtests.jl

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
using ImageSegmentation, Images, Test, SimpleWeightedGraphs, LightGraphs, StaticArrays, RegionTrees
1+
using ImageSegmentation, Images, Test, SimpleWeightedGraphs, LightGraphs, StaticArrays, DataStructures
2+
using RegionTrees: isleaf, Cell, split!
3+
using MetaGraphs: MetaGraph, clear_props!, get_prop, has_prop, set_prop!, props, vertices
24

35
using Documenter
46
doctest(ImageSegmentation, manual = false)
@@ -10,4 +12,4 @@ include("fast_scanning.jl")
1012
include("watershed.jl")
1113
include("region_merging.jl")
1214
include("meanshift.jl")
13-
15+
include("merge_segments.jl")

0 commit comments

Comments
 (0)