cv::pyrMeanShiftFiltering 和 cuda::meanShiftSegmentation源码阅读笔记

最近学了cuda，准备找个项目来练练手。先分析了一把代码，发现主要性能瓶颈在cv::pyrMeanShiftFiltering，原本打算自己撸一个gpu版的，没想到一搜发现已经有人写好了cuda::meanShiftSegmentation。那就改变计划，读读源码吧。

mean shift

mean shift简单理解就是划一个圈圈，算出mean，然后往这个方向shift过去。

图片来自opencv官方文档
思路还是挺简单、直观。但是具体到应用的时候，就可以发挥各种脑洞去改造了。

我之前用meanshift是用来聚类后自动划档，比如统计购买同一个商品的用户的注册时间，可以按新老用户划分出几个人群。因为每一个商品的用户分布都不一致，不能直接拍脑袋定一个固定值，聚成多少个簇也是不固定的。而meanshift就可以很方便解决划分多少个档位，怎么划分的问题。

那回到图像的问题，图片是rows*cols个像素，跟opencv举例子的那个图的情况好像不一样，图片怎么meanshift呢？图像有颜色空间（R,G,B）和位置空间（x,y）5维，我之前还以为是直接对颜色空间进行meanshift，其实不是的。
那先来看看cv::pyrMeanShiftFiltering是怎么处理的吧。

cv::pyrMeanShiftFiltering源码

pyr是指金字塔pyramids，那金字塔的meanshift又是怎么实现呢？

/****************************************************************************************\
*                                         Meanshift                                      *
\****************************************************************************************/

void cv::pyrMeanShiftFiltering( InputArray _src, OutputArray _dst,
                                double sp0, double sr, int max_level,
                                TermCriteria termcrit )
{
    CV_INSTRUMENT_REGION()

    Mat src0 = _src.getMat();

    if( src0.empty() )
        return;

    _dst.create( src0.size(), src0.type() );
    Mat dst0 = _dst.getMat();

    // cn=channel num?
    const int cn = 3;
    const int MAX_LEVELS = 8;

    if( (unsigned)max_level > (unsigned)MAX_LEVELS )
        CV_Error( CV_StsOutOfRange, "The number of pyramid levels is too large or negative" );

    std::vector<cv::Mat> src_pyramid(max_level+1);
    std::vector<cv::Mat> dst_pyramid(max_level+1);
    cv::Mat mask0;
    int i, j, level;
    //uchar* submask = 0;


    //(c0,c1,c2) 跟 (ofs0,ofs1,ofs2)的颜色欧式距离的平方。用平方就不用开根号了
    #define cdiff(ofs0) (tab[c0-dptr[ofs0]+255] + \
        tab[c1-dptr[(ofs0)+1]+255] + tab[c2-dptr[(ofs0)+2]+255] >= isr22)

    double sr2 = sr * sr;
    // i是指integer, isr2^2为啥最小要16？因为sr至少要2
    int isr2 = cvRound(sr2), isr22 = MAX(isr2,16);
    int tab[768];


    if( src0.type() != CV_8UC3 )
        CV_Error( CV_StsUnsupportedFormat, "Only 8-bit, 3-channel images are supported" );

    if( src0.type() != dst0.type() )
        CV_Error( CV_StsUnmatchedFormats, "The input and output images must have the same type" );

    if( src0.size() != dst0.size() )
        CV_Error( CV_StsUnmatchedSizes, "The input and output images must have the same size" );

    if( !(termcrit.type & CV_TERMCRIT_ITER) )
        termcrit.maxCount = 5;
    termcrit.maxCount = MAX(termcrit.maxCount,1);
    termcrit.maxCount = MIN(termcrit.maxCount,100);
    if( !(termcrit.type & CV_TERMCRIT_EPS) )
        termcrit.epsilon = 1.f;
    termcrit.epsilon = MAX(termcrit.epsilon, 0.f);

    // 预处理的表。tab[0] 表示 (-255)^2, tab[255]表示0^2
    // 颜色如果只有0-255，那好像512就够了，为啥要768？
    for( i = 0; i < 768; i++ )
        tab[i] = (i - 255)*(i - 255);

    // 1. construct pyramid
    src_pyramid[0] = src0;
    dst_pyramid[0] = dst0;
    for( level = 1; level <= max_level; level++ )
    {
        src_pyramid[level].create( (src_pyramid[level-1].rows+1)/2,
                        (src_pyramid[level-1].cols+1)/2, src_pyramid[level-1].type() );
        dst_pyramid[level].create( src_pyramid[level].rows,
                        src_pyramid[level].cols, src_pyramid[level].type() );
        cv::pyrDown( src_pyramid[level-1], src_pyramid[level], src_pyramid[level].size() );
        //CV_CALL( cvResize( src_pyramid[level-1], src_pyramid[level], CV_INTER_AREA ));
    }

    // 申请空间，可以复用
    mask0.create(src0.rows, src0.cols, CV_8UC1);
    //CV_CALL( submask = (uchar*)cvAlloc( (sp+2)*(sp+2) ));

    // 2. apply meanshift, starting from the pyramid top (i.e. the smallest layer)
    for( level = max_level; level >= 0; level-- )
    {
        cv::Mat src = src_pyramid[level];
        cv::Size size = src.size();
        const uchar* sptr = src.ptr();
        int sstep = (int)src.step;
        uchar* mask = 0;
        int mstep = 0;
        uchar* dptr;
        int dstep;
        float sp = (float)(sp0 / (1 << level));
        sp = MAX( sp, 1 );

        if( level < max_level )
        {
            cv::Size size1 = dst_pyramid[level+1].size();
            cv::Mat m( size.height, size.width, CV_8UC1, mask0.ptr() );
            //step是指到图片下一行要跳过多少字节
            dstep = (int)dst_pyramid[level+1].step;
            // 因为下面要算一个点到周围8个点的距离，所以边缘一个像素那一圈要跳过
            dptr = dst_pyramid[level+1].ptr() + dstep + cn;
            mstep = (int)m.step;
            mask = m.ptr() + mstep;
            //cvResize( dst_pyramid[level+1], dst_pyramid[level], CV_INTER_CUBIC );
            cv::pyrUp( dst_pyramid[level+1], dst_pyramid[level], dst_pyramid[level].size() );
            m.setTo(cv::Scalar::all(0));

            // mask的尺寸只有当前金字塔dptr的1/2，所以mask在x或者y方向都要跳两倍
            // 也就是rows方向要加mstep*2, 下面columns方向[j*2 -1]
            // 最后dilate就是消除尺寸的影响。
            // 那跟直接生成一个小的mask，再线性插值扩大一倍的区别？
            // 线性插值1的边缘可能生成一些0.8之类的灰度值。而delate可以描边扩充成1
            // 而mask这里其实需要是二值的
            for( i = 1; i < size1.height-1; i++, dptr += dstep - (size1.width-2)*3, mask += mstep*2 )
            {
                for( j = 1; j < size1.width-1; j++, dptr += cn )
                {
                    int c0 = dptr[0], c1 = dptr[1], c2 = dptr[2];
                    // 计算周围8个点是否有一个颜色距离过大
                    // * * *
                    // * 0 *
                    // * * *
                    mask[j*2 - 1] = cdiff(-3) || cdiff(3) || cdiff(-dstep-3) || cdiff(-dstep) ||
                        cdiff(-dstep+3) || cdiff(dstep-3) || cdiff(dstep) || cdiff(dstep+3);
                }
            }
            // kernel是空的Mat，那就是默认的3*3
            cv::dilate( m, m, cv::Mat() );
            mask = m.ptr();
        }
        // 金字塔的作用就是生成一个mask
        //

        dptr = dst_pyramid[level].ptr();
        dstep = (int)dst_pyramid[level].step;

        for( i = 0; i < size.height; i++, sptr += sstep - size.width*3,
                                          dptr += dstep - size.width*3,
                                          mask += mstep )
        {
            for( j = 0; j < size.width; j++, sptr += 3, dptr += 3 )
            {
                int x0 = j, y0 = i, x1, y1, iter;
                int c0, c1, c2;

                // 最开始mask是0，所以也不会越界
                // 碰到边界就停止，可以保持图像边缘
                if( mask && !mask[j] )
                    continue;

                c0 = sptr[0], c1 = sptr[1], c2 = sptr[2];

                // iterate meanshift procedure
                for( iter = 0; iter < termcrit.maxCount; iter++ )
                {
                    const uchar* ptr;
                    int x, y, count = 0;
                    int minx, miny, maxx, maxy;
                    int s0 = 0, s1 = 0, s2 = 0, sx = 0, sy = 0;
                    double icount;
                    int stop_flag;

                    // 注意这里mean shift实现时候的处理细节
                    // 先按空间划出一个方框（而不是圆形），再只对这个框里颜色相近的像素（过滤这些像素之后就有点像举例子的那个图了）进行xy求平均，算出要shift的方向。比较简洁。或者说为了实现方便。
                    // 反正idea还是比较直接的，就是把一个点追踪到它老家。
                    //mean shift: process pixels in window (p-sigmaSp)x(p+sigmaSp)
                    minx = cvRound(x0 - sp); minx = MAX(minx, 0);
                    miny = cvRound(y0 - sp); miny = MAX(miny, 0);
                    maxx = cvRound(x0 + sp); maxx = MIN(maxx, size.width-1);
                    maxy = cvRound(y0 + sp); maxy = MIN(maxy, size.height-1);
                    ptr = sptr + (miny - i)*sstep + (minx - j)*3;

                    for( y = miny; y <= maxy; y++, ptr += sstep - (maxx-minx+1)*3 )
                    {
                        int row_count = 0;
                        x = minx;
                        #if CV_ENABLE_UNROLLED
                        for( ; x + 3 <= maxx; x += 4, ptr += 12 )
                        {
                            int t0 = ptr[0], t1 = ptr[1], t2 = ptr[2];
                            if( tab[t0-c0+255] + tab[t1-c1+255] + tab[t2-c2+255] <= isr2 )
                            {
                                s0 += t0; s1 += t1; s2 += t2;
                                sx += x; row_count++;
                            }
                            t0 = ptr[3], t1 = ptr[4], t2 = ptr[5];
                            if( tab[t0-c0+255] + tab[t1-c1+255] + tab[t2-c2+255] <= isr2 )
                            {
                                s0 += t0; s1 += t1; s2 += t2;
                                sx += x+1; row_count++;
                            }
                            t0 = ptr[6], t1 = ptr[7], t2 = ptr[8];
                            if( tab[t0-c0+255] + tab[t1-c1+255] + tab[t2-c2+255] <= isr2 )
                            {
                                s0 += t0; s1 += t1; s2 += t2;
                                sx += x+2; row_count++;
                            }
                            t0 = ptr[9], t1 = ptr[10], t2 = ptr[11];
                            if( tab[t0-c0+255] + tab[t1-c1+255] + tab[t2-c2+255] <= isr2 )
                            {
                                s0 += t0; s1 += t1; s2 += t2;
                                sx += x+3; row_count++;
                            }
                        }
                        // 这个宏没有else的。因为是移动ptr，不加速移就慢慢移呗
                        // unrolled是cpu加速用的，一次性读取连续内存
                        #endif
                        for( ; x <= maxx; x++, ptr += 3 )
                        {
                            int t0 = ptr[0], t1 = ptr[1], t2 = ptr[2];
                            if( tab[t0-c0+255] + tab[t1-c1+255] + tab[t2-c2+255] <= isr2 )
                            {
                                s0 += t0; s1 += t1; s2 += t2;
                                sx += x; row_count++;
                            }
                        }
                        count += row_count;
                        sy += y*row_count;
                    }

                    if( count == 0 )
                        break;

                    icount = 1./count;
                    x1 = cvRound(sx*icount);
                    y1 = cvRound(sy*icount);
                    s0 = cvRound(s0*icount);
                    s1 = cvRound(s1*icount);
                    s2 = cvRound(s2*icount);

                    stop_flag = (x0 == x1 && y0 == y1) || std::abs(x1-x0) + std::abs(y1-y0) +
                        tab[s0 - c0 + 255] + tab[s1 - c1 + 255] +
                        tab[s2 - c2 + 255] <= termcrit.epsilon;

                    x0 = x1; y0 = y1;
                    c0 = s0; c1 = s1; c2 = s2;

                    if( stop_flag )
                        break;
                }

                dptr[0] = (uchar)c0;
                dptr[1] = (uchar)c1;
                dptr[2] = (uchar)c2;
            }
        }
    }
}

看完源码，对meanshift具体实现就比较清楚了。回到最初的问题，金字塔的作用是啥？我们可以发现金字塔的作用只是生成一个mask，这个mask是从小到大计算，每次计算meanshift先看一下小一倍的mask提取出那些图像边缘，碰到这些边缘就停止。那加了金字塔之后，可以保留多个维度的边缘，那就是meanshift出来的图像块比较小，那就是色彩层次会更“丰富”。

但如果图像的渐变比较丰富，实际看上去会比较糊，像被高斯模糊了一样。
因为cuda版本并没有这个金字塔的功能，不过还好看上去对效果影响也没有特别大，去掉问题也不大。

cuda::meanShiftSegmentation

meanshift针对每个像素点，计算互相不依赖，很适合gpu并行。

void cv::cuda::meanShiftSegmentation(InputArray _src, OutputArray _dst, int sp, int sr, int minsize, TermCriteria criteria, Stream& stream)
{
    GpuMat src = _src.getGpuMat();

    CV_Assert( src.type() == CV_8UC4 );

    const int nrows = src.rows;
    const int ncols = src.cols;
    const int hr = sr;
    const int hsp = sp;

    // Perform mean shift procedure and obtain region and spatial maps
    // 跟cpu版不同，除了保存了颜色(d_rmap），还保存最后meanshift的x,y位置(d_spmap)
    GpuMat d_rmap, d_spmap;
    cuda::meanShiftProc(src, d_rmap, d_spmap, sp, sr, criteria, stream);

    stream.waitForCompletion();

    Mat rmap(d_rmap);
    Mat spmap(d_spmap);

    Graph<SegmLinkVal> g(nrows * ncols, 4 * (nrows - 1) * (ncols - 1)
                                        + (nrows - 1) + (ncols - 1));
    // 后面还有一大坨代码，看得我挺蛋疼的。

标准的meanshift还是比较容易看懂的，后续又搞了一大坨东西。比如djset是disjoint set(并查集)，搞acm的应该有印象吧。而且还弄了个graph。
其实就是每个像素点生成下面这样4条边。然后最右边那一条再补充往下的边，最底下的一条需要补充往右的边。

其实就是增加了空间约束的floodfill吧，最后还把小的component给合并了。所以gpu版的不但不支持pyr，还反过来合并了色块。

速度对比

cpu1540 msecs.
gpu 83 msecs