path: root/ml/dlib/dlib/image_transforms/segment_image.h

// Copyright (C) 2011  Davis E. King (davis@dlib.net)
// License: Boost Software License   See LICENSE.txt for the full license.
#ifndef DLIB_SEGMENT_ImAGE_Hh_
#define DLIB_SEGMENT_ImAGE_Hh_

#include "segment_image_abstract.h"
#include "../algs.h"
#include <vector>
#include "../geometry.h"
#include "../disjoint_subsets.h"
#include "../set.h"

namespace dlib
{

// ----------------------------------------------------------------------------------------

    namespace impl
    {
        template <typename T>
        inline T edge_diff_uint(
            const T& a,
            const T& b
        )
        {
            if (a > b)
                return a - b;
            else
                return b - a;
        }

    // ----------------------------------------

        template <typename T, typename enabled = void>
        struct edge_diff_funct 
        {
            typedef double diff_type;

            template <typename pixel_type>
            double operator()(
                const pixel_type& a,
                const pixel_type& b
            ) const
            {
                return length(pixel_to_vector<double>(a) - pixel_to_vector<double>(b));
            }
        };

        template <>
        struct edge_diff_funct<uint8,void>
        { 
            typedef uint8 diff_type; 
            uint8 operator()( const uint8& a, const uint8& b) const { return edge_diff_uint(a,b); } 
        };

        template <>
        struct edge_diff_funct<uint16,void>
        { 
            typedef uint16 diff_type; 
            uint16 operator()( const uint16& a, const uint16& b) const { return edge_diff_uint(a,b); } 
        };

        template <>
        struct edge_diff_funct<uint32,void>
        { 
            typedef uint32 diff_type; 
            uint32 operator()( const uint32& a, const uint32& b) const { return edge_diff_uint(a,b); } 
        };

        template <>
        struct edge_diff_funct<double,void>
        { 
            typedef double diff_type; 
            double operator()( const double& a, const double& b) const { return std::abs(a-b); } 
        };

        template <typename T>
        struct edge_diff_funct<T, typename enable_if<is_matrix<T> >::type>
        {
            typedef double diff_type;
            double operator()(
                const T& a,
                const T& b
            ) const
            {
                return length(a-b);
            }
        };

    // ------------------------------------------------------------------------------------

        template <typename T>
        struct graph_image_segmentation_data_T
        {
            graph_image_segmentation_data_T() : component_size(1), internal_diff(0) {}
            unsigned long component_size;
            T internal_diff;
        };

    // ------------------------------------------------------------------------------------

        template <typename T>
        struct segment_image_edge_data_T
        {
            segment_image_edge_data_T (){}

            segment_image_edge_data_T (
                const rectangle& rect,
                const point& p1,
                const point& p2,
                const T& diff_
            ) :
                idx1(p1.y()*rect.width() + p1.x()),
                idx2(p2.y()*rect.width() + p2.x()),
                diff(diff_)
            {}

            bool operator<(const segment_image_edge_data_T& item) const
            { return diff < item.diff; }

            unsigned long idx1;
            unsigned long idx2;
            T diff;
        };

    // ------------------------------------------------------------------------------------

        template <typename image_view_type>
        struct uint8_or_uint16_pixels
        {
            typedef typename image_view_type::pixel_type pixel_type;
            const static bool value = is_same_type<pixel_type,uint8>::value ||
                is_same_type<pixel_type,uint16>::value;
        };

        // This is an overload of get_pixel_edges() that is optimized to segment images
        // with 8bit or 16bit  pixels very quickly.  We do this by using a radix sort
        // instead of quicksort.
        template <typename in_image_type, typename T>
        typename enable_if<uint8_or_uint16_pixels<in_image_type> >::type 
        get_pixel_edges (
            const in_image_type& in_img,
            std::vector<segment_image_edge_data_T<T> >& sorted_edges
        )
        {
            typedef typename in_image_type::pixel_type ptype;
            typedef T diff_type;
            std::vector<unsigned long> counts(std::numeric_limits<ptype>::max()+1, 0);

            edge_diff_funct<ptype> edge_diff;

            border_enumerator be(get_rect(in_img), 1);
            // we are going to do a radix sort on the edge weights.  So the first step
            // is to accumulate them into count.
            const rectangle area = get_rect(in_img);
            while (be.move_next())
            {
                const long r = be.element().y();
                const long c = be.element().x();
                const ptype pix = in_img[r][c];
                if (area.contains(c-1,r))   counts[edge_diff(pix, in_img[r  ][c-1])] += 1;
                if (area.contains(c+1,r))   counts[edge_diff(pix, in_img[r  ][c+1])] += 1;
                if (area.contains(c  ,r-1)) counts[edge_diff(pix, in_img[r-1][c  ])] += 1;
                if (area.contains(c  ,r+1)) counts[edge_diff(pix, in_img[r+1][c  ])] += 1;
            }
            for (long r = 1; r+1 < in_img.nr(); ++r)
            {
                for (long c = 1; c+1 < in_img.nc(); ++c)
                {
                    const ptype pix = in_img[r][c];
                    counts[edge_diff(pix, in_img[r-1][c+1])] += 1;
                    counts[edge_diff(pix, in_img[r  ][c+1])] += 1;
                    counts[edge_diff(pix, in_img[r+1][c  ])] += 1;
                    counts[edge_diff(pix, in_img[r+1][c+1])] += 1;
                }
            }

            const unsigned long num_edges = shrink_rect(area,1).area()*4 + in_img.nr()*2*3 - 4 + (in_img.nc()-2)*2*3;
            typedef segment_image_edge_data_T<T> segment_image_edge_data;
            sorted_edges.resize(num_edges);

            // integrate counts.  The idea is to have sorted_edges[counts[i]] be the location that edges
            // with an edge_diff of i go.  So counts[0] == 0, counts[1] == number of 0 edge diff edges, etc.
            unsigned long prev = counts[0];
            for (unsigned long i = 1; i < counts.size(); ++i)
            {
                const unsigned long temp = counts[i];
                counts[i] += counts[i-1];
                counts[i-1] -= prev;
                prev = temp;
            }
            counts[counts.size()-1] -= prev;


            // now build a sorted list of all the edges
            be.reset();
            while(be.move_next())
            {
                const point p = be.element();
                const long r = p.y();
                const long c = p.x();
                const ptype pix = in_img[r][c];
                if (area.contains(c-1,r))
                {
                    const diff_type diff = edge_diff(pix, in_img[r  ][c-1]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c-1,r),diff);
                }

                if (area.contains(c+1,r))
                {
                    const diff_type diff = edge_diff(pix, in_img[r  ][c+1]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r),diff);
                }

                if (area.contains(c  ,r-1))
                {
                    const diff_type diff = edge_diff(pix, in_img[r-1][c  ]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r-1),diff);
                }

                if (area.contains(c  ,r+1))
                {
                    const diff_type diff = edge_diff(pix, in_img[r+1][c  ]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r+1),diff);
                }
            }
            // same thing as the above loop but now we do it on the interior of the image and therefore
            // don't have to include the boundary checking if statements used above.
            for (long r = 1; r+1 < in_img.nr(); ++r)
            {
                for (long c = 1; c+1 < in_img.nc(); ++c)
                {
                    const point p(c,r);
                    const ptype pix = in_img[r][c];
                    diff_type diff;

                    diff = edge_diff(pix, in_img[r  ][c+1]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r),diff);
                    diff = edge_diff(pix, in_img[r-1][c+1]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r-1),diff);
                    diff = edge_diff(pix, in_img[r+1][c+1]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r+1),diff);
                    diff = edge_diff(pix, in_img[r+1][c  ]);
                    sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r+1),diff);
                }
            }
        }
        
    // ----------------------------------------------------------------------------------------

        // This is the general purpose version of get_pixel_edges().  It handles all pixel types.
        template <typename in_image_type, typename T>
        typename disable_if<uint8_or_uint16_pixels<in_image_type> >::type 
        get_pixel_edges (
            const in_image_type& in_img,
            std::vector<segment_image_edge_data_T<T> >& sorted_edges
        )
        {   
            const rectangle area = get_rect(in_img);
            sorted_edges.reserve(area.area()*4);

            typedef typename in_image_type::pixel_type ptype;
            edge_diff_funct<ptype> edge_diff;
            typedef T diff_type;
            typedef segment_image_edge_data_T<T> segment_image_edge_data;

            border_enumerator be(get_rect(in_img), 1);

            // now build a sorted list of all the edges
            be.reset();
            while(be.move_next())
            {
                const point p = be.element();
                const long r = p.y();
                const long c = p.x();
                const ptype& pix = in_img[r][c];
                if (area.contains(c-1,r))
                {
                    const diff_type diff = edge_diff(pix, in_img[r  ][c-1]);
                    sorted_edges.push_back(segment_image_edge_data(area,p,point(c-1,r),diff));
                }

                if (area.contains(c+1,r))
                {
                    const diff_type diff = edge_diff(pix, in_img[r  ][c+1]);
                    sorted_edges.push_back(segment_image_edge_data(area,p,point(c+1,r),diff));
                }

                if (area.contains(c  ,r-1))
                {
                    const diff_type diff = edge_diff(pix, in_img[r-1][c  ]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r-1),diff));
                }
                if (area.contains(c  ,r+1))
                {
                    const diff_type diff = edge_diff(pix, in_img[r+1][c  ]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r+1),diff));
                }
            }
            // same thing as the above loop but now we do it on the interior of the image and therefore
            // don't have to include the boundary checking if statements used above.
            for (long r = 1; r+1 < in_img.nr(); ++r)
            {
                for (long c = 1; c+1 < in_img.nc(); ++c)
                {
                    const point p(c,r);
                    const ptype& pix = in_img[r][c];
                    diff_type diff;

                    diff = edge_diff(pix, in_img[r  ][c+1]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r),diff));
                    diff = edge_diff(pix, in_img[r+1][c+1]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r+1),diff));
                    diff = edge_diff(pix, in_img[r+1][c  ]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r+1),diff));
                    diff = edge_diff(pix, in_img[r-1][c+1]);
                    sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r-1),diff));
                }
            }

            std::sort(sorted_edges.begin(), sorted_edges.end());

        }

    // ------------------------------------------------------------------------------------

    } // end of namespace impl

// ----------------------------------------------------------------------------------------

    template <
        typename in_image_type,
        typename out_image_type
        >
    void segment_image (
        const in_image_type& in_img_,
        out_image_type& out_img_,
        const double k = 200,
        const unsigned long min_size = 10
    )
    {
        using namespace dlib::impl;
        typedef typename image_traits<in_image_type>::pixel_type ptype;
        typedef typename edge_diff_funct<ptype>::diff_type diff_type;

        // make sure requires clause is not broken
        DLIB_ASSERT(is_same_object(in_img_, out_img_) == false,
            "\t void segment_image()"
            << "\n\t The input images can't be the same object."
            );

        COMPILE_TIME_ASSERT(is_unsigned_type<typename image_traits<out_image_type>::pixel_type>::value);

        const_image_view<in_image_type> in_img(in_img_);
        image_view<out_image_type> out_img(out_img_);

        out_img.set_size(in_img.nr(), in_img.nc());
        // don't bother doing anything if the image is too small
        if (in_img.nr() < 2 || in_img.nc() < 2)
        {
            assign_all_pixels(out_img,0);
            return;
        }

        disjoint_subsets sets;
        sets.set_size(in_img.size());

        std::vector<segment_image_edge_data_T<diff_type> > sorted_edges;
        get_pixel_edges(in_img, sorted_edges);

        std::vector<graph_image_segmentation_data_T<diff_type> > data(in_img.size());

        // now start connecting blobs together to make a minimum spanning tree.
        for (unsigned long i = 0; i < sorted_edges.size(); ++i)
        {
            const unsigned long idx1 = sorted_edges[i].idx1;
            const unsigned long idx2 = sorted_edges[i].idx2;

            unsigned long set1 = sets.find_set(idx1);
            unsigned long set2 = sets.find_set(idx2);
            if (set1 != set2)
            {
                const diff_type diff = sorted_edges[i].diff;
                const diff_type tau1 = static_cast<diff_type>(k/data[set1].component_size);
                const diff_type tau2 = static_cast<diff_type>(k/data[set2].component_size);

                const diff_type mint = std::min(data[set1].internal_diff + tau1, 
                                                data[set2].internal_diff + tau2);
                if (diff <= mint)
                {
                    const unsigned long new_set = sets.merge_sets(set1, set2);
                    data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                    data[new_set].internal_diff = diff;
                }
            }
        }

        // now merge any really small blobs
        if (min_size != 0)
        {
            for (unsigned long i = 0; i < sorted_edges.size(); ++i)
            {
                const unsigned long idx1 = sorted_edges[i].idx1;
                const unsigned long idx2 = sorted_edges[i].idx2;

                unsigned long set1 = sets.find_set(idx1);
                unsigned long set2 = sets.find_set(idx2);
                if (set1 != set2 && (data[set1].component_size < min_size || data[set2].component_size < min_size))
                {
                    const unsigned long new_set = sets.merge_sets(set1, set2);
                    data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                    //data[new_set].internal_diff = sorted_edges[i].diff;
                }
            }
        }

        unsigned long idx = 0;
        for (long r = 0; r < out_img.nr(); ++r)
        {
            for (long c = 0; c < out_img.nc(); ++c)
            {
                out_img[r][c] = sets.find_set(idx++);
            }
        }
    }

// ----------------------------------------------------------------------------------------
// ----------------------------------------------------------------------------------------
//                     Candidate object location generation code.
// ----------------------------------------------------------------------------------------
// ----------------------------------------------------------------------------------------

    namespace impl
    {
        struct edge_data
        {
            double edge_diff;
            unsigned long set1;  
            unsigned long set2;
            bool operator<(const edge_data& item) const
            {
                return edge_diff < item.edge_diff;
            }
        };

        template <
            typename in_image_type,
            typename diff_type
            >
        void find_basic_candidate_object_locations (
            const in_image_type& in_img,
            const std::vector<dlib::impl::segment_image_edge_data_T<diff_type> >& sorted_edges,
            std::vector<rectangle>& out_rects,
            std::vector<edge_data>& edges,
            const double k,
            const unsigned long min_size 
        )
        {
            using namespace dlib::impl;

            std::vector<dlib::impl::segment_image_edge_data_T<diff_type> > rejected_edges;
            rejected_edges.reserve(sorted_edges.size());

            out_rects.clear();
            edges.clear();

            // don't bother doing anything if the image is too small
            if (in_img.nr() < 2 || in_img.nc() < 2)
            {
                return;
            }

            disjoint_subsets sets;
            sets.set_size(in_img.size());


            std::vector<graph_image_segmentation_data_T<diff_type> > data(in_img.size());



            std::pair<unsigned long,unsigned long> last_blob_edge(std::numeric_limits<unsigned long>::max(),
                                                                  std::numeric_limits<unsigned long>::max());;
            // now start connecting blobs together to make a minimum spanning tree.
            for (unsigned long i = 0; i < sorted_edges.size(); ++i)
            {
                const unsigned long idx1 = sorted_edges[i].idx1;
                const unsigned long idx2 = sorted_edges[i].idx2;

                unsigned long set1 = sets.find_set(idx1);
                unsigned long set2 = sets.find_set(idx2);
                if (set1 != set2)
                {
                    const diff_type diff = sorted_edges[i].diff;
                    const diff_type tau1 = static_cast<diff_type>(k/data[set1].component_size);
                    const diff_type tau2 = static_cast<diff_type>(k/data[set2].component_size);

                    const diff_type mint = std::min(data[set1].internal_diff + tau1, 
                        data[set2].internal_diff + tau2);
                    if (diff <= mint)
                    {
                        const unsigned long new_set = sets.merge_sets(set1, set2);
                        data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                        data[new_set].internal_diff = diff;
                    }
                    else
                    {
                        // Don't bother keeping multiple edges from the same pair of blobs, we
                        // only need one for what we will do later.
                        if (std::make_pair(set1,set2) != last_blob_edge)
                        {
                            segment_image_edge_data_T<diff_type> temp = sorted_edges[i];
                            temp.idx1 = set1;
                            temp.idx2 = set2;
                            rejected_edges.push_back(temp);
                            last_blob_edge = std::make_pair(set1,set2);
                        }
                    }
                }
            }


            // merge small blobs
            for (unsigned long i = 0; i < rejected_edges.size(); ++i)
            {
                const unsigned long idx1 = rejected_edges[i].idx1;
                const unsigned long idx2 = rejected_edges[i].idx2;

                unsigned long set1 = sets.find_set(idx1);
                unsigned long set2 = sets.find_set(idx2);
                rejected_edges[i].idx1 = set1;
                rejected_edges[i].idx2 = set2;
                if (set1 != set2 && (data[set1].component_size < min_size || data[set2].component_size < min_size))
                {
                    const unsigned long new_set = sets.merge_sets(set1, set2);
                    data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                    data[new_set].internal_diff = rejected_edges[i].diff;
                }
            }

            // find bounding boxes of each blob
            std::map<unsigned long, rectangle> boxes;
            std::map<unsigned long, unsigned long> box_id_map;
            unsigned long idx = 0;
            for (long r = 0; r < in_img.nr(); ++r)
            {
                for (long c = 0; c < in_img.nc(); ++c)
                {
                    const unsigned long id = sets.find_set(idx++);
                    // Accumulate the current point into its box and if it is the first point
                    // in the box then also record the id number for this box.
                    if ((boxes[id] += point(c,r)).area() == 1)
                        box_id_map[id] = boxes.size()-1;
                }
            }

            // copy boxes into out_rects
            out_rects.resize(boxes.size());
            for (std::map<unsigned long,rectangle>::iterator i = boxes.begin(); i != boxes.end(); ++i)
            {
                out_rects[box_id_map[i->first]] = i->second;
            }

            // Now find the edges between the boxes 
            typedef dlib::memory_manager<char>::kernel_2c mm_type;
            dlib::set<std::pair<unsigned long, unsigned long>, mm_type>::kernel_1a neighbors_final;
            for (unsigned long i = 0; i < rejected_edges.size(); ++i)
            {
                const unsigned long idx1 = rejected_edges[i].idx1;
                const unsigned long idx2 = rejected_edges[i].idx2;

                unsigned long set1 = sets.find_set(idx1);
                unsigned long set2 = sets.find_set(idx2);
                if (set1 != set2)
                {
                    std::pair<unsigned long, unsigned long> p = std::make_pair(set1,set2);
                    if (!neighbors_final.is_member(p))
                    {
                        neighbors_final.add(p);

                        edge_data temp;
                        const diff_type mint = std::min(data[set1].internal_diff , 
                                                        data[set2].internal_diff );
                        temp.edge_diff = rejected_edges[i].diff - mint;
                        temp.set1 = box_id_map[set1];
                        temp.set2 = box_id_map[set2];
                        edges.push_back(temp);
                    }
                }
            }

            std::sort(edges.begin(), edges.end());
        }
    } // end namespace impl

// ----------------------------------------------------------------------------------------

    template <typename alloc>
    void remove_duplicates (
        std::vector<rectangle,alloc>& rects
    )
    {
        std::sort(rects.begin(), rects.end(), std::less<rectangle>());
        unsigned long num_unique = 1;
        for (unsigned long i = 1; i < rects.size(); ++i)
        {
            if (rects[i] != rects[i-1])
            {
                rects[num_unique++] = rects[i];
            }
        }
        if (rects.size() != 0)
            rects.resize(num_unique);
    }

// ----------------------------------------------------------------------------------------

    template <
        typename in_image_type,
        typename EXP
        >
    void find_candidate_object_locations (
        const in_image_type& in_img_,
        std::vector<rectangle>& rects,
        const matrix_exp<EXP>& kvals,
        const unsigned long min_size = 20,
        const unsigned long max_merging_iterations = 50
    )
    {
        // make sure requires clause is not broken
        DLIB_ASSERT(is_vector(kvals) && kvals.size() > 0,
            "\t void find_candidate_object_locations()"
            << "\n\t Invalid inputs were given to this function."
            << "\n\t is_vector(kvals): " << is_vector(kvals)
            << "\n\t kvals.size():     " << kvals.size()
            );

        typedef dlib::memory_manager<char>::kernel_2c mm_type;
        typedef dlib::set<rectangle, mm_type>::kernel_1a set_of_rects;

        using namespace dlib::impl;
        typedef typename image_traits<in_image_type>::pixel_type ptype;
        typedef typename edge_diff_funct<ptype>::diff_type diff_type;

        const_image_view<in_image_type> in_img(in_img_);

        // don't bother doing anything if the image is too small
        if (in_img.nr() < 2 || in_img.nc() < 2)
        {
            return;
        }

        std::vector<edge_data> edges;
        std::vector<rectangle> working_rects;
        std::vector<segment_image_edge_data_T<diff_type> > sorted_edges;
        get_pixel_edges(in_img, sorted_edges);

        disjoint_subsets sets;

        for (long j = 0; j < kvals.size(); ++j)
        {
            const double k = kvals(j);

            find_basic_candidate_object_locations(in_img, sorted_edges, working_rects, edges, k, min_size);
            rects.insert(rects.end(), working_rects.begin(), working_rects.end());


            // Now iteratively merge all the rectangles we have and record the results.
            // Note that, unlike what is described in the paper 
            //    Segmentation as Selective Search for Object Recognition" by Koen E. A. van de Sande, et al.
            // we don't use any kind of histogram/SIFT like thing to order the edges
            // between the blobs.  Here we simply order by the pixel difference value.
            // Additionally, note that we keep progressively merging boxes in the outer
            // loop rather than performing just a single iteration as indicated in the
            // paper.
            set_of_rects detected_rects;
            bool did_merge = true;
            for (unsigned long iter = 0; did_merge && iter < max_merging_iterations; ++iter) 
            {
                did_merge = false;
                sets.clear();
                sets.set_size(working_rects.size());

                // recursively merge neighboring blobs until we have merged everything
                for (unsigned long i = 0; i < edges.size(); ++i)
                {
                    edge_data temp = edges[i];

                    temp.set1 = sets.find_set(temp.set1);
                    temp.set2 = sets.find_set(temp.set2);
                    if (temp.set1 != temp.set2)
                    {
                        rectangle merged_rect = working_rects[temp.set1] + working_rects[temp.set2];
                        // Skip merging this pair of blobs if it was merged in a previous
                        // iteration.  Doing this lets us consider other possible blob
                        // merges.
                        if (!detected_rects.is_member(merged_rect))
                        {
                            const unsigned long new_set = sets.merge_sets(temp.set1, temp.set2);
                            rects.push_back(merged_rect);
                            working_rects[new_set] = merged_rect;
                            did_merge = true;
                            detected_rects.add(merged_rect);
                        }
                    }
                }
            }
        }

        remove_duplicates(rects);
    }

// ----------------------------------------------------------------------------------------

    template <
        typename in_image_type
        >
    void find_candidate_object_locations (
        const in_image_type& in_img,
        std::vector<rectangle>& rects
    )
    {
        find_candidate_object_locations(in_img, rects, linspace(50, 200, 3));
    }

// ----------------------------------------------------------------------------------------

}

#endif // DLIB_SEGMENT_ImAGE_Hh_