implementation of `merge_hierarchical` in Julia

graph_merge.jl

import Colors
using DataStructures
using ImageSegmentation
using LightGraphs
using MetaGraphs

# idea from  https://vcansimplify.wordpress.com/2014/08/17/hierarchical-merging-of-region-adjacency-graphs/

function _weight_mean_color(data::MetaGraph, n1::Int, n2::Int)::Real
    # Determine weight of an edge between n1 and n2
    return Colors.colordiff(
        get_prop(data, n1, :mean_color), 
        get_prop(data, n2, :mean_color)
    )
end


function merge_nodes(g::MetaGraph, e::AbstractEdge)
    src, dst = e.src, e.dst
    clr = get_prop(g, dst, :total_color) + get_prop(g, src, :total_color)
    npx = get_prop(g, dst, :pixel_count) + get_prop(g, src, :pixel_count)

    set_prop!(g, dst, :total_color, clr)
    set_prop!(g, dst, :pixel_count, npx)
    set_prop!(g, dst, :mean_color, clr / npx)
    set_prop!(g, dst, :labels, vcat(
         get_prop(g, src, :labels),
         get_prop(g, dst, :labels)
    ))

    # Clear props on the now unused node src, to make its obsolescence clear.
    clear_props!(g, src)
end


function make_new_edges(g::MetaGraph, e::AbstractEdge)::Array{Edge}
    edges = []
    edge_neighbors = union(Set(neighbors(g, e.src)), Set(neighbors(g, e.dst)))
    for n in setdiff(edge_neighbors, e.src, e.dst) 
        edge = Edge(e.dst, n)
        add_edge!(g, edge)
        set_prop!(g, edge, :weight, _weight_mean_color(g, e.dst, n))
        push!(edges, edge)
    end

    return edges
end


function invalidate_neighbors(g::MetaGraph, edge_heap::MutableBinaryHeap, e::AbstractEdge)
    """Invalidates the neighbors of e in the edge_heap. """
    function invalidate(src, dst)
        for n in setdiff(Set(neighbors(g, src)), dst)
            edge = Edge(src, n)
            h = get_prop(g, edge, :handle)
            update!(edge_heap, h, (edge, false))
        end
    end
    invalidate(e.src, e.dst)
    invalidate(e.dst, e.src)
end


function seg_to_graph(seg::SegmentedImage)::MetaGraph
    weight(i, j) = Colors.colordiff(segment_mean(seg, i), segment_mean(seg, j))
    rag, _ = region_adjacency_graph(seg, weight)

    g = MetaGraph(rag)
    for v in vertices(rag)
        set_prop!(g, v, :labels, [v])
        set_prop!(g, v, :pixel_count, seg.segment_pixel_count[v])
        set_prop!(g, v, :mean_color, seg.segment_means[v])
        set_prop!(g, v, :total_color, seg.segment_means[v] * seg.segment_pixel_count[v])
    end

    for e in edges(rag)
        set_prop!(g, Edge(e.src, e.dst), :weight, e.weight)
    end
    return g
end


function merge_hierarchical(seg::SegmentedImage, thresh::Number)::SegmentedImage
    g = seg_to_graph(seg)

    function weight(t::Tuple{Edge{Int64}, Bool})::Real
        return has_prop(g, t[1], :weight) ?  get_prop(g, t[1], :weight) : 0
    end
    edge_heap = MutableBinaryHeap{Tuple{Edge{Int64}, Bool}}(Base.By(weight), [])
    sizehint!(edge_heap, 3 * length(edges(g)))  # Overkill, or not enough?
    for e in edges(g)
        handle = push!(edge_heap, (e, true))
        set_prop!(g, e, :handle, handle)
    end

    while !isempty(edge_heap) && weight(first(edge_heap)) < thresh
        e, valid = pop!(edge_heap)
        if valid
            # Invalidate all edges touching this edge (heap removal would be expensive)
            invalidate_neighbors(g, edge_heap, e)

            # Merge the two nodes into one (this involves merging their props)
            merge_nodes(g, e)

            # Make new edges to the merged node
            new_edges = make_new_edges(g, e)

            # Remove edges to src. 
            # don't call rem_vertex!; it renumbers all vertices, causing chaos.
            for n in collect(neighbors(g, e.src))
                rem_edge!(g, e.src, n)
            end

            # Add new edges to heap.
            for e in new_edges
                handle = push!(edge_heap, (e, true))
                set_prop!(g, e, :handle, handle)
            end
        end
    end

    return resegment(seg, g)
end


function merge_simple(seg::SegmentedImage, thresh::Real)::SegmentedImage
    """Simpler implementation of the same idea. No heap, runtime likely worse for large graphs"""
    g = seg_to_graph(seg)

    function min_edge(g, f)
        es = collect(edges(g))
        w, ix = findmin(map(f, es))
        e = es[ix]
        return w, e
    end

    weight = e -> get_prop(g, e, :weight)
    w, e = min_edge(g, weight)
    while w < thresh
        merge_nodes(g, e)
        make_new_edges(g, e)

        for n in collect(neighbors(g, e.src))
            rem_edge!(g, e.src, n)
        end

        w, e = min_edge(g, weight)
    end

    return resegment(seg, g)
end


function resegment(seg::SegmentedImage, g::MetaGraph)::SegmentedImage
    """Takes a segmentation and a region adjacency graph produced by merge_hierarchical

    Produces a segmentation that corresponds to the graph.
    """
    remaining = collect(filter(v -> 0 < length(props(g, v)), vertices(g)))

    px_labels = copy(seg.image_indexmap)
    # Need to to do this in two passes to avoid re-using labels.
    # first merge labels
    for v in remaining
        labels = get_prop(g, v, :labels)
        for l in labels
            if l != v
                ix = findall(x -> x == l, px_labels)
                px_labels[ix] .= v
            end
        end
    end 

    means, px_counts = Dict{Int64, Colorant}(), Dict{Int64, Int64}() 
    for (i, v) in enumerate(remaining)
        ix = findall(x -> x == v, px_labels)
        px_labels[ix] .= i
        means[i] = get_prop(g, v, :mean_color)
        px_counts[i] = get_prop(g, v, :pixel_count)
    end

    labels = collect(1:length(remaining))

    return SegmentedImage(px_labels, labels, means, px_counts)
end