|
| 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 | + |
0 commit comments