Add earth mover's distance

Author: timholy

src/ImageDistances.jl

@@ -1,5 +1,3 @@
-__precompile__()
-
 module ImageDistances
 
 import Base: size, ==
@@ -9,10 +7,12 @@ import Distances: evaluate, colwise, pairwise
 
 using Colors
 using ProgressMeter
+using EarthMoversDistance
 
 include("generic.jl")
 include("hausdorff.jl")
 include("ciede2000.jl")
+include("earthmover.jl")
 
 export
     # generic types
@@ -24,6 +24,7 @@ export
     Hausdorff,
     ModifiedHausdorff,
     CIEDE2000,
+    EarthMoverBinned,
 
     # helper functions
     hausdorff,

src/earthmover.jl

@@ -0,0 +1,67 @@
+using EarthMoversDistance: CSignature, CFlow, FLOW_ARRAY_SIZE
+
+struct EarthMoverBinned{F,H<:AbstractArray,C} <: ImageMetric
+    binner!::F
+    histA::H
+    histB::H
+    sigA::CSignature
+    sigB::CSignature
+    cost::C
+    cflow::Vector{CFlow}
+end
+
+function costlb(ip, jp, edges::AbstractVector{T}) where T
+    # lower-bound cost
+    # Adjacent bins have no cost, because you may move by "epsilon" and cross over
+    # to the next bin
+    i, j = ip[], jp[]
+    ii, jj = Int(i), Int(j)
+    # The final bin is for NaN pixels, which match with zero cost
+    (abs(ii-jj) <= 1 || ii == length(edges)+2 || jj == length(edges)+2) && return zero(Cfloat)
+    return Cfloat(ii > jj ? edges[ii-1] - edges[jj] : edges[jj-1] - edges[ii])
+end
+
+function EarthMoverBinned(edges::AbstractVector, cost=(i, j)->costlb(i, j, edges))
+    binner! = graybinner(edges)
+    len = length(edges)+2
+    len > FLOW_ARRAY_SIZE && error("edges exceeded maximum allowed size of $(EarthMoversDistance.FLOW_ARRAY_SIZE)")
+    histA, histB = Vector{Cfloat}(undef, len), Vector{Cfloat}(undef, len)
+    sigA, sigB = CSignature(histA), CSignature(histB)
+    costfun = @cfunction $cost Cfloat (Ref{Cfloat}, Ref{Cfloat})
+    cflow = fill(CFlow(0, 0, 0), FLOW_ARRAY_SIZE)
+    metric = EarthMoverBinned(binner!, histA, histB, sigA, sigB, costfun, cflow)
+    # For reasons that are not clear, the first usage errors unless we do the following:
+    evaluate(metric, [0.0f0], [0.0f0])
+    return metric
+end
+
+function evaluate(em::EarthMoverBinned, imgA::AbstractArray, imgB::AbstractArray)
+    em.binner!(em.histA, imgA)
+    em.binner!(em.histB, imgB)
+    # Do the ccall directly to avoid the allocations
+    cflowsizeptr = Ref{Cint}(0)
+    res = ccall((:emd, :emd), # function name and library name
+                Cfloat,       # return type
+                (Ref{CSignature}, Ref{CSignature}, Ptr{Nothing}, Ref{CFlow}, Ref{Cint}), # argument types
+                Ref(em.sigA), Ref(em.sigB), em.cost, em.cflow, cflowsizeptr)
+    return res
+end
+
+function graybinner(edges::AbstractVector)
+    @assert !Base.has_offset_axes(edges)
+    ord = Base.Order.ForwardOrdering()
+    @assert issorted(edges, ord)
+    let ord=ord   # julia issue #15276
+        function graybinner!(hist, A)
+            fill!(hist, 0)
+            for a in A
+                if isnan(a)
+                    hist[end] += oneunit(eltype(hist))
+                else
+                    hist[searchsortedfirst(edges, a, ord)] += oneunit(eltype(hist))
+                end
+            end
+            return hist
+        end
+    end
+end

test/runtests.jl

@@ -53,4 +53,25 @@ evaluate(d::MyImgDist, imgA, imgB) = π
       @test ciede2000(A, B) ≥ 0
       @test ciede2000(A, B) == ciede2000(B, A)
     end
+
+    @testset "Earth Mover" begin
+        edges = 0.0:0.2:1.0
+        metric = EarthMoverBinned(edges)
+        imgA = fill(0.1, 5, 5)
+        imgB = fill(0.3, 5, 5)
+        @test evaluate(metric, imgA, imgA) == 0
+        @test evaluate(metric, imgB, imgB) == 0
+        # Since A and B are on opposite sides of the bin boundary, the cost should be 0
+        # (i.e., epsilon)
+        @test evaluate(metric, imgA, imgB) < 1e-8
+        imgB = fill(0.5, 5, 5)
+        @test evaluate(metric, imgA, imgB) ≈ step(edges)
+        imgB = [fill(0.1, 5) fill(0.3, 5) fill(0.5, 5) fill(0.7, 5) fill(0.9, 5)]
+        @test evaluate(metric, imgA, imgB) ≈ (5 + 10 + 15) * step(edges) / 25
+        # Images don't have to be the same size
+        imgB = [fill(0.1, 10) fill(0.3, 10) fill(0.5, 10) fill(0.7, 10) fill(0.9, 10)]
+        @test evaluate(metric, imgA, imgB) ≈ 5 * step(edges) / 25
+        imgB = [fill(0.9, 10) fill(0.3, 10) fill(0.7, 10) fill(0.7, 10) fill(0.9, 10)]
+        @test evaluate(metric, imgA, imgB) ≈ 15*(2*step(edges))/25
+    end
 end