Image segmentation - Split and Merge (Quadtrees)

Is there an implementation for the split and merge method of image segmentation? any advice would be muc开发者_JS百科h appreciated.

What segmentation is about?

Segmentation means division of your image into several connected regions. Basically, you could do segmentation with two definitions of region: you could define a region as a group of connected similar pixels, or a set of connected pixels surrounded by discontinuities (edges). Split and merge uses the first approach.

Mathematically speaking: if your whole image is represented by the set of pixels (called R), than you would like to get subsets such as

1. Segmentation is complete, so all subregions sum up to the whole R. Union of all regions is R₁ U R₂ U ... U R_n = R.
2. R_i is connected.
3. Regions are distinct. R_i∩R_j=∅ given that i≠j
4. Regions have similar properties. This could be expressed by a function called the homogenity criterion (P). It should give TRUE for the members of the given region, and FALSE for all the other regions.
5. Neighbor regions cannot been merged. For all regions P(R_i U R_j)=FALSE given that i≠j.

What split and merge algorithm is about?

So first, we have to choose a homogenity criterion. A homogenity criterion could be global (depending on the whole region) or local (depending on a small window of the region, and if it's true for all windows, than it's true for the region). A simple example could be the deviation from average should be smaller than a threshold. ∀p_i∈R_i: |p_i-μ|≤f*σ.

The split and merge algorithm have two phases: the split, and the merge phase. In the split phase we recursively split regions into four subregions (starting with the whole image as one region) until our homogenity criterion is met in all subregions. It's easy to see that the 1-4 conditions of segmentation are met. We proceed to merge step in order to satisfy the 5^th condition.

In the merge step we check that P(R_i U R_j)=TRUE for each two neighbor regions, and merge the two regions. We repeat this step until no more changes are necessary. Now we met all the conditions, we had our image segmented into subregions.

Pseudocode

Here is a pseudocode to split and merge algorithm:

1. Init: we have only one big region (the whole image).
2. Split: If P(R_i)=TRUE proceed to next step. Otherwise subdivide R_i to four subregions and perform step 2 on them.
3. Merge: If R_i and R_j are neighbors and P(R_i U R_j) = TRUE, merge the two regions, than repeat step 3. If there are no such regions we are finished.

This is my implementation. I am not a c++ / opencv guru so if someone find out some way to optimize this script add comments please!

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <iostream>

using namespace cv;
using namespace std;

Mat img;
Size size;

struct region {
    // tree data structure
    vector<region> childs;
    bool validity; // TODO: have a method for clear the data structure and remove regions with false validity

    // tree for split&merge procedure
    Rect roi;
    Mat m;
    Scalar label;
    Mat mask; // for debug. don't use in real cases because it is computationally too heavy.
};

//----------------------------------------------------------------------------------------------------------------------- merging
bool mergeTwoRegion(region& parent, const Mat& src, region& r1, region& r2,  bool (*predicate)(const Mat&)) {
    if(r1.childs.size()==0 && r2.childs.size()==0) {

        Rect roi1 = r1.roi;
        Rect roi2 = r2.roi;
        Rect roi12 = roi1 | roi2;
        if(predicate( src(roi12) )) {
            r1.roi = roi12;

            // recompute mask
            r1.mask = Mat::zeros(size, CV_8U);
            rectangle(r1.mask, r1.roi, 1, CV_FILLED);

            r2.validity = false;
            return true;
        }
    }
    return false;
}

void merge(const Mat& src, region& r, bool (*predicate)(const Mat&)) {
    // check for adjiacent regions. if predicate is true, then  merge.
    // the problem is to check for adjiacent regions.. one way can be:
    // check merging for rows. if neither rows can be merged.. check for cols.

    bool row1=false, row2=false, col1=false, col2=false;

    if(r.childs.size()<1) return;

    // try with the row
    row1 = mergeTwoRegion(r, src, r.childs[0], r.childs[1], predicate);
    row2 = mergeTwoRegion(r, src, r.childs[2], r.childs[3], predicate);

    if( !(row1 | row2) ) {
        // try with column
        col1 = mergeTwoRegion(r, src, r.childs[0], r.childs[2], predicate);
        col2 = mergeTwoRegion(r, src, r.childs[1], r.childs[3], predicate);
    } 

    for(int i=0; i<r.childs.size(); i++) {
        if(r.childs[i].childs.size()>0) 
            merge(src, r.childs[i], predicate);
    }
}

//----------------------------------------------------------------------------------------------------------------------- quadtree splitting
region split(const Mat& src, Rect roi, bool (*predicate)(const Mat&)) {
    vector<region> childs;
    region r;

    r.roi = roi;
    r.m = src;
    r.mask = Mat::zeros(size, CV_8U);
    rectangle(r.mask, r.roi, 1, CV_FILLED);
    r.validity = true;

    bool b = predicate(src);
    if(b) {
        Scalar mean, s;
        meanStdDev(src, mean, s);
        r.label = mean;
    } else {
        int w = src.cols/2;
        int h = src.rows/2;
        region r1 = split(src(Rect(0,0, w,h)), Rect(roi.x, roi.y, w,h), predicate);
        region r2 = split(src(Rect(w,0, w,h)), Rect(roi.x+w, roi.y, w,h), predicate);
        region r3 = split(src(Rect(0,h, w,h)), Rect(roi.x, roi.y+h, w,h), predicate);
        region r4 = split(src(Rect(w,h, w,h)), Rect(roi.x+w, roi.y+h, w,h), predicate);
        r.childs.push_back( r1 );
        r.childs.push_back( r2 );
        r.childs.push_back( r3 );
        r.childs.push_back( r4 );
    }
    //merge(img, r, predicate);
    return r;
}

//----------------------------------------------------------------------------------------------------------------------- tree traversing utility
void print_region(region r) {
    if(r.validity==true && r.childs.size()==0) {
        cout << r.mask << " at " << r.roi.x << "-" << r.roi.y << endl;
        cout << r.childs.size() << endl;
        cout << "---" << endl;
    }
    for(int i=0; i<r.childs.size(); i++) {
        print_region(r.childs[i]);
    }
}

void draw_rect(Mat& imgRect, region r) {
    if(r.validity==true && r.childs.size()==0) 
        rectangle(imgRect, r.roi, 50, .1);
    for(int i=0; i<r.childs.size(); i++) {
        draw_rect(imgRect, r.childs[i]);
    }
}

void draw_region(Mat& img, region r) {
    if(r.validity==true && r.childs.size()==0) 
        rectangle(img, r.roi, r.label, CV_FILLED);
    for(int i=0; i<r.childs.size(); i++) {
        draw_region(img, r.childs[i]);
    }
}

//----------------------------------------------------------------------------------------------------------------------- split&merge test predicates
bool predicateStdZero(const Mat& src) {
    Scalar stddev, mean;
    meanStdDev(src, mean, stddev);
    return stddev[0]==0;
}
bool predicateStd5(const Mat& src) {
    Scalar stddev, mean;
    meanStdDev(src, mean, stddev);
    return (stddev[0]<=5.8) || (src.rows*src.cols<=25);
}

//----------------------------------------------------------------------------------------------------------------------- main
int main( int /*argc*/, char** /*argv*/ )
{
    img = (Mat_<uchar>(4,4) << 0,0,1,1,
                               1,1,1,1, 
                               3,3,3,3,
                               3,4,4,3);

    cout << img << endl;
    size = img.size();

    region r;
    r = split(img, Rect(0,0,img.cols,img.rows), &predicateStdZero);
    merge(img, r, &predicateStdZero);
    cout << "------- print" << endl;
    print_region(r);

    cout << "-----------------------" << endl;

    img = imread("lena.jpg", 0);

    // round (down) to the nearest power of 2 .. quadtree dimension is a pow of 2.
    int exponent = log(min(img.cols, img.rows)) / log (2);
    int s = pow(2.0, (double)exponent);
    Rect square = Rect(0,0, s,s);
    img = img(square).clone();

    namedWindow("original", CV_WINDOW_AUTOSIZE);
    imshow( "original", img );

    cout << "now try to split.." << endl;
    r = split(img, Rect(0,0,img.cols,img.rows), predicateStd5);

    cout << "splitted" << endl;
    Mat imgRect = img.clone();
    draw_rect(imgRect, r);
    namedWindow("split", CV_WINDOW_AUTOSIZE);
    imshow( "split", imgRect );
    imwrite("split.jpg", imgRect);

    merge(img, r, &predicateStd5);
    Mat imgMerge = img.clone();
    draw_rect(imgMerge, r);
    namedWindow("merge", CV_WINDOW_AUTOSIZE);
    imshow( "merge", imgMerge );
    imwrite( "merge.jpg", imgMerge );

    Mat imgSegmented = img.clone();
    draw_region(imgSegmented, r);
    namedWindow("segmented", CV_WINDOW_AUTOSIZE);
    imshow( "segmented", imgSegmented );
    imwrite( "segmented.jpg", imgSegmented );

    while( true )
    {
        char c = (char)waitKey(10);
        if( c == 27 ) { break; }
    }

    return 0;
}