samples/cpp/grabcut.cpp

一个使用GrabCut算法的示例

#include "opencv2/imgcodecs.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgproc.hpp"
 
#include <iostream>
 
using namespace std;
using namespace cv;
 
static void help(char** argv)
{
cout << "\n本程序演示了GrabCut图像分割——在一个区域中选择一个对象\n"
            "然后GrabCut会尝试将其分割出来。\n"
            "调用:\n"
<< argv[0] << " <图像名>\n"
            "\n选择您想要分割的对象周围的矩形区域\n" <<
            "\n热键：\n"
            "\tESC - 退出程序\n"
            "\tr - 恢复原始图像\n"
            "\tn - 下一次迭代\n"
            "\n"
            "\t鼠标左键 - 设置矩形\n"
            "\n"
            "\tCTRL+鼠标左键 - 设置GC_BGD像素\n"
            "\tSHIFT+鼠标左键 - 设置GC_FGD像素\n"
            "\n"
            "\tCTRL+鼠标右键 - 设置GC_PR_BGD像素\n"
            "\tSHIFT+鼠标右键 - 设置GC_PR_FGD像素\n" << endl;
}
 
const Scalar RED = Scalar(0,0,255);
const Scalar PINK = Scalar(230,130,255);
const Scalar BLUE = Scalar(255,0,0);
const Scalar LIGHTBLUE = Scalar(255,255,160);
const Scalar GREEN = Scalar(0,255,0);
 
const int BGD_KEY = EVENT_FLAG_CTRLKEY;
const int FGD_KEY = EVENT_FLAG_SHIFTKEY;
 
static void getBinMask( const Mat& comMask, Mat& binMask )
{
    if( comMask.empty() || comMask.type()!=CV_8UC1 )
        CV_Error( Error::StsBadArg, "comMask为空或类型不正确（非CV_8UC1）" );
    if( binMask.empty() || binMask.rows!=comMask.rows || binMask.cols!=comMask.cols )
binMask.create( comMask.size(), CV_8UC1 );
binMask = comMask & 1;
}
 
class GCApplication
{
public:
    enum{ NOT_SET = 0, IN_PROCESS = 1, SET = 2 };
    static const int radius = 2;
    static const int thickness = -1;
 
    void reset();
    void setImageAndWinName( const Mat& _image, const string& _winName );
    void showImage() const;
    void mouseClick( int event, int x, int y, int flags, void* param );
    int nextIter();
    int getIterCount() const { return iterCount; }
private:
    void setRectInMask();
    void setLblsInMask( int flags, Point p, bool isPr );
 
    const string* winName;
    const Mat* image;
    Mat mask;
    Mat bgdModel, fgdModel;
 
    uchar rectState, lblsState, prLblsState;
    bool isInitialized;
 
    Rect rect;
vector<Point> fgdPxls, bgdPxls, prFgdPxls, prBgdPxls;
    int iterCount;
};
 
void GCApplication::reset()
{
    if( !mask.empty() )
        mask.setTo(Scalar::all(GC_BGD));
bgdPxls.clear(); fgdPxls.clear();
prBgdPxls.clear(); prFgdPxls.clear();
 
isInitialized = false;
rectState = NOT_SET;
lblsState = NOT_SET;
prLblsState = NOT_SET;
iterCount = 0;
}
 
void GCApplication::setImageAndWinName( const Mat& _image, const string& _winName )
{
    if( _image.empty() || _winName.empty() )
        return;
image = &_image;
winName = &_winName;
    mask.create( image->size(), CV_8UC1);
reset();
}
 
void GCApplication::showImage() const
{
    if( image->empty() || winName->empty() )
        return;
 
    Mat res;
    Mat binMask;
image->copyTo( res );
    if( isInitialized ){
getBinMask( mask, binMask);
 
        Mat black (binMask.rows, binMask.cols, CV_8UC3, cv::Scalar(0,0,0));
black.setTo(Scalar::all(255), binMask);
 
        addWeighted(black, 0.5, res, 0.5, 0.0, res);
    }
 
vector<Point>::const_iterator it;
    for( it = bgdPxls.begin(); it != bgdPxls.end(); ++it )
        circle( res, *it, radius, BLUE, thickness );
    for( it = fgdPxls.begin(); it != fgdPxls.end(); ++it )
        circle( res, *it, radius, RED, thickness );
    for( it = prBgdPxls.begin(); it != prBgdPxls.end(); ++it )
        circle( res, *it, radius, LIGHTBLUE, thickness );
    for( it = prFgdPxls.begin(); it != prFgdPxls.end(); ++it )
        circle( res, *it, radius, PINK, thickness );
 
    if( rectState == IN_PROCESS || rectState == SET )
        rectangle( res, Point( rect.x, rect.y ), Point(rect.x + rect.width, rect.y + rect.height ), GREEN, 2);
 
    imshow( *winName, res );
}
 
void GCApplication::setRectInMask()
{
    CV_Assert( !mask.empty() );
    mask.setTo( GC_BGD );
rect.x = max(0, rect.x);
rect.y = max(0, rect.y);
rect.width = min(rect.width, image->cols-rect.x);
rect.height = min(rect.height, image->rows-rect.y);
(mask(rect)).setTo( Scalar(GC_PR_FGD) );
}
 
void GCApplication::setLblsInMask( int flags, Point p, bool isPr )
{
vector<Point> *bpxls, *fpxls;
    uchar bvalue, fvalue;
    if( !isPr )
    {
bpxls = &bgdPxls;
fpxls = &fgdPxls;
bvalue = GC_BGD;
fvalue = GC_FGD;
    }
    else
    {
bpxls = &prBgdPxls;
fpxls = &prFgdPxls;
bvalue = GC_PR_BGD;
fvalue = GC_PR_FGD;
    }
    if( flags & BGD_KEY )
    {
bpxls->push_back(p);
        circle( mask, p, radius, bvalue, thickness );
    }
    if( flags & FGD_KEY )
    {
fpxls->push_back(p);
        circle( mask, p, radius, fvalue, thickness );
    }
}
 
void GCApplication::mouseClick( int event, int x, int y, int flags, void* )
{
    // TODO 添加不良参数检查
    switch( event )
    {
    case EVENT_LBUTTONDOWN: // 设置矩形或GC_BGD(GC_FGD)标签
        {
            bool isb = (flags & BGD_KEY) != 0,
isf = (flags & FGD_KEY) != 0;
            if( rectState == NOT_SET && !isb && !isf )
            {
rectState = IN_PROCESS;
rect = Rect( x, y, 1, 1 );
            }
            if ( (isb || isf) && rectState == SET )
lblsState = IN_PROCESS;
        }
        break;
    case EVENT_RBUTTONDOWN: // 设置GC_PR_BGD(GC_PR_FGD)标签
        {
            bool isb = (flags & BGD_KEY) != 0,
isf = (flags & FGD_KEY) != 0;
            if ( (isb || isf) && rectState == SET )
prLblsState = IN_PROCESS;
        }
        break;
    case EVENT_LBUTTONUP
        if( rectState == IN_PROCESS )
        {
            if(rect.x == x || rect.y == y){
rectState = NOT_SET;
            }
            else{
rect = Point( rect.x, rect.y ), Point(x,y) );
rectState = SET;
setRectInMask();
                CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
            }
            showImage();
        }
        if( lblsState == IN_PROCESS )
        {
setLblsInMask(flags, Point(x,y), false);
lblsState = SET;
nextIter();
            showImage();
        }
        else{
            if(rectState == SET){
nextIter();
                showImage();
            }
        }
        break;
    case EVENT_RBUTTONUP
        if( prLblsState == IN_PROCESS )
        {
setLblsInMask(flags, Point(x,y), true);
prLblsState = SET;
        }
        if(rectState == SET){
nextIter();
            showImage();
        }
        break;
    case EVENT_MOUSEMOVE
        if( rectState == IN_PROCESS )
        {
rect = Point( rect.x, rect.y ), Point(x,y) );
            CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
            showImage();
        }
        else if( lblsState == IN_PROCESS )
        {
setLblsInMask(flags, Point(x,y), false);
            showImage();
        }
        else if( prLblsState == IN_PROCESS )
        {
setLblsInMask(flags, Point(x,y), true);
            showImage();
        }
        break;
    }
}
 
int GCApplication::nextIter()
{
    if( isInitialized )
        grabCut( *image, mask, rect, bgdModel, fgdModel, 1 );
    else
    {
        if( rectState != SET )
            return iterCount;
 
        if( lblsState == SET || prLblsState == SET )
            grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_MASK );
        else
            grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_RECT );
 
isInitialized = true;
    }
iterCount++;
 
bgdPxls.clear(); fgdPxls.clear();
prBgdPxls.clear(); prFgdPxls.clear();
 
    return iterCount;
}
 
GCApplication gcapp;
 
static void on_mouse( int event, int x, int y, int flags, void* param )
{
gcapp.mouseClick( event, x, y, flags, param );
}
 
int main( int argc, char** argv )
{
    cv::CommandLineParser parser(argc, argv, "{@input| messi5.jpg |}");
help(argv);
 
    string filename = parser.get<string>("@input");
    if( filename.empty() )
    {
cout << "\n糟糕，文件名为空" << endl;
        return 1;
    }
    Mat image = imread(samples::findFile(filename), IMREAD_COLOR);
    if( image.empty() )
    {
cout << "\n 糟糕，无法读取图像文件：" << filename << endl;
        return 1;
    }
 
    const string winName = "image";
    namedWindow( winName, WINDOW_AUTOSIZE );
    setMouseCallback( winName, on_mouse, 0 );
 
gcapp.setImageAndWinName( image, winName );
gcapp.showImage();
 
    for(;;)
    {
        char c = (char)waitKey(0);
        switch( c )
        {
        case '\x1b'
cout << "正在退出..." << endl;
            goto exit_main;
        case 'r'
cout << endl;
gcapp.reset();
gcapp.showImage();
            break;
        case 'n'
            int iterCount = gcapp.getIterCount();
cout << "<" << iterCount << "... ";
            int newIterCount = gcapp.nextIter();
            if( newIterCount > iterCount )
            {
gcapp.showImage();
cout << iterCount << ">" << endl;
            }
            else
cout << "必须确定矩形>" << endl;
            break;
        }
    }
 
exit_main
    destroyWindow( winName );
    return 0;
}