distance_transform.h

#ifndef CVD_INC_DISTANCE_TRANSFORM_HPP
#define CVD_INC_DISTANCE_TRANSFORM_HPP

#include <cmath>
#include <cvd/vision_exceptions.h>

#include <algorithm>
#include <cmath>
#include <cvd/image.h>
#include <limits>
#include <vector>

namespace CVD
{

template <class Precision = double>
class DistanceTransformEuclidean
{
    // Implements
    // Distance Transforms of Sampled Functions
    //   Pedro F. Felzenszwalb Daniel P. Huttenlocher

    public:
    DistanceTransformEuclidean()
        : sz(ImageRef(-1, -1))
        , big_number(1'000'000'000) //Hmm, why doesn't HUGE_VAL work?
    //Anyway, hilariously small number here so it works with int, too.
    {
    }

    private:
    inline void transform_row(const int n)
    {
        v[0] = 0; //std::numeric_limits<Precision>::infinity();
        z[0] = -big_number; //std::numeric_limits<Precision>::infinity();
        z[1] = +big_number; // std::numeric_limits<Precision>::infinity();
        int k = 0;
        for(int q = 1; q < n; q++)
        {
            Precision s = ((f[q] + (q * q)) - (f[v[k]] + (v[k] * v[k]))) / (2 * q - 2 * v[k]);
            while(s <= z[k])
            {
                k--;
                s = ((f[q] + (q * q)) - (f[v[k]] + (v[k] * v[k]))) / (2 * q - 2 * v[k]);
            }
            k++;
            v[k] = q;
            z[k] = s;
            z[k + 1] = +big_number; //std::numeric_limits<Precision>::infinity();
        }
        k = 0;
        for(int q = 0; q < n; q++)
        {
            while(z[k + 1] < q)
            {
                k++;
            }
            d[q] = ((q - v[k]) * (q - v[k])) + f[v[k]];
            pos[q] = q > v[k] ? (q - v[k]) : (v[k] - q);
            //cout << "n=" << n << " q=" << q << " d[q]=" << d[q] << " v[k]=" << v[k] << " f[v[k]]=" << f[v[k]] << endl;
        }
    }

    void transform_image(BasicImage<Precision>& DT)
    {
        const ImageRef img_sz(DT.size());
        for(int x = 0; x < img_sz.x; x++)
        {
            for(int y = 0; y < img_sz.y; y++)
            {
                f[y] = DT[y][x];
            }
            transform_row(img_sz.y);
            for(int y = 0; y < img_sz.y; y++)
            {
                DT[y][x] = d[y];
            }
        }
        //#if 0
        for(int y = 0; y < img_sz.y; y++)
        {
            for(int x = 0; x < img_sz.x; x++)
            {
                f[x] = DT[y][x];
            }
            transform_row(img_sz.x);
            for(int x = 0; x < img_sz.x; x++)
            {
                DT[y][x] = d[x];
            }
        }
        //#endif
    }

    template <class Functor>
    void transform_image_with_ADT(BasicImage<Precision>& DT, BasicImage<ImageRef>& ADT, const Functor& func)
    {
        const ImageRef img_sz(DT.size());
        const double maxdist = img_sz.x * img_sz.y;
        for(int x = 0; x < img_sz.x; x++)
        {
            for(int y = 0; y < img_sz.y; y++)
            {
                f[y] = DT[y][x];
            }
            transform_row(img_sz.y);
            for(int y = 0; y < img_sz.y; y++)
            {
                DT[y][x] = d[y];
            }
        }
        for(int y = 0; y < img_sz.y; y++)
        {
            for(int x = 0; x < img_sz.x; x++)
            {
                f[x] = DT[y][x];
            }
            transform_row(img_sz.x);
            for(int x = 0; x < img_sz.x; x++)
            {
                DT[y][x] = d[x];
                const double hyp = d[x];
                if(hyp >= maxdist)
                {
                    ADT[y][x] = ImageRef(-1, -1);
                    continue;
                }
                const double dx = pos[x];
                const double dy = sqrt(hyp - dx * dx);
                const int ddy = static_cast<int>(dy);
                const ImageRef candA(static_cast<int>(x - dx), static_cast<int>(y - ddy));
                const ImageRef candB(static_cast<int>(x - dx), static_cast<int>(y + ddy));
                const ImageRef candC(static_cast<int>(x + dx), static_cast<int>(y - ddy));
                const ImageRef candD(static_cast<int>(x + dx), static_cast<int>(y + ddy));
                if(DT.in_image(candA) && func(candA))
                {
                    ADT[y][x] = candA;
                }
                else if(DT.in_image(candB) && func(candB))
                {
                    ADT[y][x] = candB;
                }
                else if(DT.in_image(candC) && func(candC))
                {
                    ADT[y][x] = candC;
                }
                else if(DT.in_image(candD) && func(candD))
                {
                    ADT[y][x] = candD;
                }
                else
                {
                    //ADT[y][x] = ImageRef(-1,-1);
                    throw Exceptions::Vision::BadInput("DistanceTransformEuclidean: no points set");
                }
            }
        }
    }

    void resize(const ImageRef& new_size)
    {
        if(new_size != sz)
        {
            sz = new_size;
            int m(std::max(new_size.x, new_size.y));
            d.resize(m);
            v.resize(m);
            f.resize(m);
            pos.resize(m);
            z.resize(m + 1);
        }
    }

    public:
    template <class T>
    void transform_ADT(const BasicImage<T>& feature, BasicImage<Precision>& DT, BasicImage<ImageRef>& ADT)
    {
        if(feature.size() != DT.size())
            throw Exceptions::Vision::IncompatibleImageSizes(__FUNCTION__);

        if(feature.size() != ADT.size())
            throw Exceptions::Vision::IncompatibleImageSizes(__FUNCTION__);

        resize(feature.size());

        apply_functor(DT, NotZero<T>(feature));

        transform_image_with_ADT(DT, ADT, NotZero<T>(feature));
    }

    void transform(BasicImage<Precision>& out)
    {
        resize(out.size());
        transform_image(out);
    }

    template <class C>
    struct NotZero
    {
        NotZero(const BasicImage<C>& im_)
            : im(im_)
        {
        }

        const BasicImage<C>& im;
        bool operator()(const ImageRef& i) const
        {
            return im[i] != 0;
        }
    };

    template <class Out, class Functor>
    void apply_functor(BasicImage<Out>& out, const Functor& f)
    {
        for(int y = 0; y < out.size().y; y++)
            for(int x = 0; x < out.size().x; x++)
                if(f(ImageRef(x, y)))
                    out[y][x] = 0;
                else
                    out[y][x] = big_number;
    }

    private:
    ImageRef sz;
    int big_number;
    std::vector<Precision> d;
    std::vector<int> v;
    std::vector<Precision> z;
    std::vector<Precision> f;
    std::vector<Precision> pos;
};

template <class T, class Q>
void euclidean_distance_transform_sq(const BasicImage<T>& in, BasicImage<Q>& out)
{
    if(in.size() != out.size())
        throw Exceptions::Vision::IncompatibleImageSizes(__FUNCTION__);
    DistanceTransformEuclidean<Q> dt;
    dt.apply_functor(out, typename DistanceTransformEuclidean<Q>::template NotZero<T>(in));
    dt.transform(out);
}

template <class T, class Q>
void euclidean_distance_transform_sq(const BasicImage<T>& in, BasicImage<Q>& out, BasicImage<ImageRef>& lookup_DT)
{
    if(in.size() != out.size())
        throw Exceptions::Vision::IncompatibleImageSizes(__FUNCTION__);
    DistanceTransformEuclidean<Q> dt;
    dt.transform_ADT(in, out, lookup_DT);
}

template <class T, class Q>
void euclidean_distance_transform(const BasicImage<T>& in, BasicImage<Q>& out)
{
    euclidean_distance_transform_sq(in, out);
    for(int y = 0; y < out.size().y; y++)
        for(int x = 0; x < out.size().x; x++)
            out[y][x] = sqrt(out[y][x]);
}

template <class T, class Q>
void euclidean_distance_transform(const BasicImage<T>& in, BasicImage<Q>& out, BasicImage<ImageRef>& lookup_DT)
{
    euclidean_distance_transform_sq(in, out, lookup_DT);
    for(int y = 0; y < out.size().y; y++)
        for(int x = 0; x < out.size().x; x++)
            out[y][x] = sqrt(out[y][x]);
}

#ifndef DOXYGEN_IGNORE_INTERNAL
namespace Internal
{
    template <class C>
    class DoDistanceTransform
    {
    };

    template <class T>
    struct ImagePromise<DoDistanceTransform<T>>
    {
        ImagePromise(const BasicImage<T>& in_)
            : in(in_)
        {
        }

        const BasicImage<T>& in;

        template <class C>
        void execute(Image<C>& im)
        {
            im.resize(in.size());
            euclidean_distance_transform(in, im);
        }
    };
};

template <class T>
Internal::ImagePromise<Internal::DoDistanceTransform<T>> euclidean_distance_transform(const BasicImage<T>& in)
{
    using namespace Internal;
    return ImagePromise<DoDistanceTransform<T>>(in);
}
#else

template <class T>
Image euclidean_distance_transform(const BasicImage<T>& in);
#endif

};
#endif