grabcut 源代码

以下是 OpenCV 中的 GrabCut 源代码，你可以在 GitHub 上找到它：

void cv::grabCut( InputArray _img, InputOutputArray _mask, Rect rect,
                  InputOutputArray _bgdModel, InputOutputArray _fgdModel,
                  int iterCount, int mode )
{
    Mat img = _img.getMat(), mask0 = _mask.getMat();
    if( img.empty() )
        CV_Error( CV_StsBadArg, "image is empty" );
    if( mask0.empty() )
        CV_Error( CV_StsBadArg, "mask is empty" );
    if( rect.width <= 0 || rect.height <= 0 )
        CV_Error( CV_StsBadArg, "rectangle width and height must be positive" );
    if( rect.x < 0 || rect.y < 0 || rect.x + rect.width > img.cols ||
        rect.y + rect.height > img.rows )
        CV_Error( CV_StsBadArg, "rectangle out of image" );

    if( iterCount <= 0 )
        return;

    int i, j, k, n;
    int x, y;
    int width = rect.width, height = rect.height, kCompCount = 5;
    int compIdx, newMaskVal, maxIter = iterCount;
    int* mask = 0;
    float* bgdData = 0;
    float* fgdData = 0;
    uchar* ptr;
    uchar* ptr2;
    uchar** maskTab = 0;
    float** bgdTab = 0;
    float** fgdTab = 0;
    int* binIdx = 0;
    double t;

    _mask.create( img.size(), CV_8UC1 );
    mask = _mask.getMat().data;

    // initialize state
    maskTab = (uchar**)cvAlloc( height*sizeof(maskTab[0]) +
                                height*width*sizeof(maskTab[0][0]) +
                                kCompCount*(sizeof(bgdTab[0]) + sizeof(fgdTab[0])) );
    bgdTab = (float**)(maskTab + height);
    fgdTab = bgdTab + kCompCount;
    bgdData = (float*)(fgdTab + kCompCount);
    fgdData = bgdData + kCompCount*3;
    binIdx = (int*)cvAlloc(256*sizeof(binIdx[0]));
    for( i = 0; i < height; i++ )
        maskTab[i] = mask + img.cols*i;

    for( i = 0; i < 256; i++ )
        binIdx[i] = i/(256/(kCompCount-1));

    t = (double)getTickCount();
    initGMMs( img(rect), maskTab, bgdTab, fgdTab );
    calcBeta( img(rect), bgdTab, fgdTab, maskTab, beta );
    t = (double)getTickCount() - t;
    if( iterCount == 1 )
        return;

    for( int count = 0; count < maxIter; count++ )
    {
        int changeCount = 0;

        // assign GMM components
        for( i = 0; i < height; i++ )
        {
            ptr = img.ptr<uchar>(i+rect.y) + rect.x*3;
            ptr2 = maskTab[i] + rect.x;

            for( j = 0; j < width; j++, ptr += 3 )
            {
                double p[5], s = 0;

                for( k = 0; k < kCompCount; k++ )
                {
                    p[k] = pdf( ptr, bgdTab[k], fgdTab[k] );
                    s += p[k];
                }

                if( s > std::numeric_limits<double>::epsilon() )
                {
                    s = 1./s;
                    for( k = 0; k < kCompCount; k++ )
                        p[k] *= s;
                }
                else
                    p[0] = 1.f, p[1] = 0, p[2] = 0, p[3] = 0, p[4] = 0;

                double max_p = p[0];
                int max_idx = 0;

                for( k = 1; k < kCompCount; k++ )
                {
                    if( max_p < p[k] )
                    {
                        max_p = p[k];
                        max_idx = k;
                    }
                }

                ptr2[j] = (uchar)max_idx;
            }
        }

        // estimate GMM parameters
        for( k = 0; k < kCompCount; k++ )
        {
            memset( bgdData + k*3, 0, 3*sizeof(float) );
            memset( fgdData + k*3, 0, 3*sizeof(float) );
            n = 0;

            for( i = 0; i < height; i++ )
            {
                ptr = img.ptr<uchar>(i+rect.y) + rect.x*3;
                ptr2 = maskTab[i] + rect.x;

                for( j = 0; j < width; j++, ptr += 3 )
                {
                    if( ptr2[j] == k )
                    {
                        bgdData[k*3 + 0] += ptr[0];
                        bgdData[k*3 + 1] += ptr[1];
                        bgdData[k*3 + 2] += ptr[2];
                        n++;
                    }
                    else if( ptr2[j] == k + kCompCount )
                    {
                        fgdData[k*3 + 0] += ptr[0];
                        fgdData[k*3 + 1] += ptr[1];
                        fgdData[k*3 + 2] += ptr[2];
                        n++;
                    }
                }
            }

            if( n != 0 )
            {
                float inv_n = 1.f/n;
                bgdData[k*3 + 0] *= inv_n; bgdData[k*3 + 1] *= inv_n;
                bgdData[k*3 + 2] *= inv_n;
                fgdData[k*3 + 0] *= inv_n; fgdData[k*3 + 1] *= inv_n;
                fgdData[k*3 + 2] *= inv_n;
            }
            else
            {
                bgdData[k*3 + 0] = float( rand() );
                bgdData[k*3 + 1] = float( rand() );
                bgdData[k*3 + 2] = float( rand() );
                fgdData[k*3 + 0] = float( rand() );
                fgdData[k*3 + 1] = float( rand() );
                fgdData[k*3 + 2] = float( rand() );
            }

            if( !calcInverseCovAndDeterm( bgdTab[k], bgdGMM[k].cov.inv(), bgdGMM[k].covDeterm ) ||
                !calcInverseCovAndDeterm( fgdTab[k], fgdGMM[k].cov.inv(), fgdGMM[k].covDeterm ) )
            {
                // if DET(COV) is too small (singular matrix), exit
                changeCount = maxIter + 1;
                break;
            }

            bgdGMM[k].weight = (float)n/(width*height);
            fgdGMM[k].weight = 1.f - bgdGMM[k].weight;
        }

        if( changeCount > maxIter )
            break;

        // assign labels
        for( i = 0; i < height; i++ )
        {
            ptr2 = maskTab[i] + rect.x;
            for( j = 0; j < width; j++ )
                ptr2[j] = binIdx[ptr2[j]];
        }

        // make sure that "hard" components are not split
        for( k = 0; k < kCompCount; k++ )
        {
            int sum_b = 0, count_b = 0;
            int sum_f = 0, count_f = 0;

            for( i = 0; i < height; i++ )
            {
                ptr = img.ptr<uchar>(i+rect.y) + rect.x*3;
                ptr2 = maskTab[i] + rect.x;

                for( j = 0; j < width; j++, ptr += 3 )
                {
                    if( ptr2[j] == k )
                    {
                        sum_b += ptr[0];
                        sum_b += ptr[1];
                        sum_b += ptr[2];
                        count_b++;
                    }
                    else if( ptr2[j] == k + kCompCount )
                    {
                        sum_f += ptr[0];
                        sum_f += ptr[1];
                        sum_f += ptr[2];
                        count_f++;
                    }
                }
            }

            if( count_b == 0 )
                bgdGMM[k].weight = 0.f;
            else
            {
                bgdData[k*3 + 0] = (float)(sum_b/count_b);
                bgdData[k*3 + 1] = (float)(sum_b/(count_b+1));
                bgdData[k*3 + 2] = (float)(sum_b/(count_b+2));
                bgdGMM[k].weight = 1.f/(count_b + 2);
            }

            if( count_f == 0 )
                fgdGMM[k].weight = 0.f;
            else
            {
                fgdData[k*3 + 0] = (float)(sum_f/count_f);
                fgdData[k*3 + 1] = (float)(sum_f/(count_f+1));
                fgdData[k*3 + 2] = (float)(sum_f/(count_f+2));
                fgdGMM[k].weight = 1.f/(count_f + 2);
            }
        }

        if( changeCount > maxIter )
            break;
    }

    cvFree( &maskTab );
    cvFree( &binIdx );
}