在上次用 CUDA 实现导向滤波 后, 想着导向滤波能以很小的 mask 还原高分辨率下的边缘, 能不能搞点事情出来, 当时正好在研究 Darknet 框架, 然后又看到 grabcut 算法, 用 opencv 试了下, 感觉效果有点意思, 后面想了下, 这几个可以连在一起, 先读取高分辨率的图像, 然后用降低分辨率先通过 yolov3 算出人物框 (非常稳定, 不跳, 几乎不会出现有人而找不到的情况), 再用 grabcut 算出低 mask, 然后用这个 mask 结合原图用导向滤波得到高分辨率下清晰边缘的分割图, 最后再把 CUDA 算出的结果直接丢给 UE4/Unity3D 的对应 DX11 中.
先放出 cuda 版 grabcut 的效果图 (博客园好像传不了视频, 复制知乎上的链接, 可以右键新标签面里打开视频看)
当然 opencv 本身提供的 grabcut 是用 CPU 算的, 416*416 差不多有个几桢左右, 肯定不满足要求, 所以在这个前提下, 用 Cuda 来实现整个 grabcut 整个流程.
因要集成到 UE4/Untiy3D 对应软件中, 在 Windows 本台下, 我们首先需要如下.
1 拿到 darknet 源码, C 风格代码, 说实话, 没想过能看到这种短小精湛的深度学习模型框架, 带 cuda 与 cudnn, 直接在 VS 下建一个动态链接库, 把相应代码放进来, 不多, 改动几个 Linux 下的函数, 包含编译符 GPU 与 CUDNN, 然后修改 darknet.h 这个函数, dllexport 这个文件下函数, 后期还可以改下, 这个文件下的编译符移到别的位置, pthread 头文件的引用这些, 以及可以直接传入 GPU 数据这几点. opencv4 虽然已经有 yolo3 的实现, 但是我们要 GPU 的, 后续的操作几乎全在 GPU 上.
2 我们主要是参照 opencv 里的 grabcut 实现, 为了更好的参数一些数据, 我们最好编译自己的 opencv 版本, 我是用的 opencv-4.0.0-alpha, 比较老的一个版本, 需要带 opencv_contrib, 包含 opencv_cuda 相关的模块, 主要是后期我们实现 cuda 版 grabcut 如果不好确认我们是否正常实现就可以调试进去看数据值, 看源码, 以及用 GpuMat/PtrStepSz, 这个简单封装真的很适合处理图像数据.
如上二点完成后, 我们可以简单分析下 grabcut 算法主要由那种算法构成, 如何完成一个流程, 一次大迭代主要如下, 主要是先定框, 框内为可能前景区, 框外为背景区, 二个区域分别为 Kmeans 分类, 然后二个分类的数据分别 GMM 建立模型, 最后用 GMM 算出最大流算法需要的每点的 source/sink, 建立 graph, 用 GPU 友好的 maxflow 算法 push-relabel 得到最小切, 也就是 mask, 然后把 mask 给 GMM 重新建立模型然后重复这个过程.
如上所说, 我们主要实现三种算法, 分别是 kmeans,gmm,push-relabel. 下面如无别的说明, 本机给的是在 1070 下 416*416 分辨率下的时间.
首先讲下 kmeans 的优化, kmeans 的实现比较简单, 就是一个不断重新分配中心点至到稳定的过程, 其中有个过程是把数据根据前后景分别归类 (我们假设聚合为 5 块), 最直观最简单的方法就是用原子操作.
- __global__ void updateClusterAtomic(PtrStepSz<uchar4> source, PtrStepSz<uchar> clusterIndex, kmeansI& meansbg, kmeansI& meansfg)
- {
- const int idx = blockDim.x * blockIdx.x + threadIdx.x;
- const int idy = blockDim.y * blockIdx.y + threadIdx.y;
- if (idx <source.cols && idy < source.rows)
- {
- float4 color = rgbauchar42float4(source(idy, idx));
- // 背景块, 使用 kcentersbg, 否则使用 kcentersfg
- kmeansI& kmeans = (!inRect(idx, idy)) ? meansbg : meansfg;
- int index = clusterIndex(idy, idx);
- atomicAdd(&kmeans.kcenters[index].x, color.x);
- atomicAdd(&kmeans.kcenters[index].y, color.y);
- atomicAdd(&kmeans.kcenters[index].z, color.z);
- atomicAdd(&kmeans.kcenters[index].w, color.w);
- }
- }
- updateClusterAtomic
用 Nsight 看了下时间, 花费超过 1ms 了, 我是想到所有线程最后归类时都只操作 10 块地方, 原子操作在这有点把多线程变成有限的几个能同时做事了.
因此需要换种想法, 很简单, 不用原子操作, 用一部分线程如 32*8, 遍历最个纹理数据, 然后在一个线程里 32*8 个数据再整合, 使用这步后, 时间降低 0.5ms 左右, 利用的还不是很完全, 毕竟一次 kmeans 稳定需要遍历十几次左右, 这个过程就要了 5ms 后面就不用搞了, 在这基础上, 结合共享显存, 每个点一个线程先用共享显存处理后放入对应 grid 块大小的显存中, 然后数据除 32*8 的 grid 块执行上面的操作, 这个操作包含后面统计等加起 0.12ms, 这样十几次迭代也就不到 2ms, 不过我没想到的是, 从 cudaMemcpyDeviceToHost 一个 int32 的四字节 Nsight 给的时间很低 0.001ms, 但是从上面的简隔来看, 应该有 0.1ms 了, 应该还有一些别的花费, 所以很多操作, 后续的统计计算我就是用 GPU 里一个核心来算, 也不先 download 给 CPU 算然后再 upload 上来, 但是这个操作还没想到办法避免, 毕竟需要在 CPU 上判定是否已经稳定后关闭循环.
- // 把 source 所有收集到一块 gridDim.x*gridDim.y 块数据上.
- template<int blockx, int blocky>
- __global__ void updateCluster(PtrStepSz<uchar4> source, PtrStepSz<uchar> clusterIndex, PtrStepSz<uchar> mask, float4* kencter, int* kindexs)
- {
- __shared__ float3 centers[blockx*blocky][CUDA_GRABCUT_K2];
- __shared__ int indexs[blockx*blocky][CUDA_GRABCUT_K2];
- const int idx = blockDim.x * blockIdx.x + threadIdx.x;
- const int idy = blockDim.y * blockIdx.y + threadIdx.y;
- const int threadId = threadIdx.x + threadIdx.y * blockDim.x;
- #pragma unroll CUDA_GRABCUT_K2
- for (int i = 0; i <CUDA_GRABCUT_K2; i++)
- {
- centers[threadId][i] = make_float3(0.f);
- indexs[threadId][i] = 0;
- }
- __syncthreads();
- if (idx < source.cols && idy < source.rows)
- {
- // 所有值都放入共享 centers
- int index = clusterIndex(idy, idx);
- bool bFg = checkFg(mask(idy, idx));
- int kindex = bFg ? index : (index + CUDA_GRABCUT_K);
- float4 color = rgbauchar42float4(source(idy, idx));
- centers[threadId][kindex] = make_float3(color);
- indexs[threadId][kindex] = 1;
- __syncthreads();
- // 每个线程块进行二分聚合, 每次操作都保存到前一半数组里, 直到最后保存在线程块里第一个线程上 (这块比较费时, 0.1ms)
- for (uint stride = blockDim.x*blockDim.y / 2; stride> 0; stride>>= 1)
- {
- //int tid = (threadId&(stride - 1));
- if (threadId <stride)//stride 2^n
- {
- #pragma unroll CUDA_GRABCUT_K2
- for (int i = 0; i < CUDA_GRABCUT_K2; i++)
- {
- centers[threadId][i] += centers[threadId + stride][i];
- indexs[threadId][i] += indexs[threadId + stride][i];
- }
- }
- //if (stride> 32)
- __syncthreads();
- }
- // 每块的第一个线程集合, 把共享 centers 存入临时显存块上 kencter
- if (threadIdx.x == 0 && threadIdx.y == 0)
- {
- int blockId = blockIdx.x + blockIdx.y * gridDim.x;
- #pragma unroll CUDA_GRABCUT_K2
- for (int i = 0; i <CUDA_GRABCUT_K2; i++)
- {
- int id = blockId * 2 * CUDA_GRABCUT_K + i;
- kencter[id] = make_float4(centers[0][i], 0.f);
- kindexs[id] = indexs[0][i];
- }
- }
- }
- }
- //block 1*1,threads(暂时选 32*4), 对如上 gridDim.x*gridDim.y 的数据用 blockx*blocky 个线程来处理
- template<int blockx, int blocky>
- __global__ void updateCluster(float4* kencter, int* kindexs, kmeansI& meansbg, kmeansI& meansfg, int& delta, int gridcount)
- {
- __shared__ float3 centers[blockx*blocky][CUDA_GRABCUT_K2];
- __shared__ int indexs[blockx*blocky][CUDA_GRABCUT_K2];
- const int threadId = threadIdx.x + threadIdx.y * blockDim.x;
- #pragma unroll CUDA_GRABCUT_K2
- for (int i = 0; i <CUDA_GRABCUT_K2; i++)
- {
- centers[threadId][i] = make_float3(0.f);
- indexs[threadId][i] = 0;
- }
- __syncthreads();
- //int gridcount = gridDim.x*gridDim.y;
- int blockcount = blockDim.x*blockDim.y;
- int count = gridcount / blockcount + 1;
- // 每块共享变量, 每个线程操纵自己对应那块地址, 不需要同步, 用线程块操作一个大内存
- for (int i = 0; i < count; i++)
- {
- int id = i*blockcount + threadId;
- if (id < gridcount)
- {
- #pragma unroll CUDA_GRABCUT_K2
- for (int j = 0; j < CUDA_GRABCUT_K2; j++)
- {
- int iid = id * CUDA_GRABCUT_K2 + j;
- centers[threadId][j] += make_float3(kencter[iid]);
- indexs[threadId][j] += kindexs[iid];
- }
- }
- }
- __syncthreads();
- for (uint stride = blockDim.x*blockDim.y / 2; stride> 0; stride>>= 1)
- {
- if (threadId <stride)
- {
- #pragma unroll CUDA_GRABCUT_K2
- for (int i = 0; i < CUDA_GRABCUT_K2; i++)
- {
- centers[threadId][i] += centers[threadId + stride][i];
- indexs[threadId][i] += indexs[threadId + stride][i];
- }
- }
- //if (stride> 32)
- __syncthreads();
- }
- if (threadIdx.x == 0 && threadIdx.y == 0)
- {
- int count = 0;
- // 收集所有信息, 并重新更新 cluster, 记录变更的大小
- for (int i = 0; i <CUDA_GRABCUT_K; i++)
- {
- meansfg.kcenters[i] = make_float4(centers[0][i], 0.f);
- if (indexs[0][i] != 0)
- meansfg.cluster[i] = meansfg.kcenters[i] / indexs[0][i];
- count += abs(indexs[0][i] - meansfg.length[i]);
- meansfg.length[i] = indexs[0][i];
- }
- for (int i = CUDA_GRABCUT_K; i < CUDA_GRABCUT_K2; i++)
- {
- meansbg.kcenters[i - CUDA_GRABCUT_K] = make_float4(centers[0][i], 0.f);
- if (indexs[0][i] != 0)
- meansbg.cluster[i - CUDA_GRABCUT_K] = meansbg.kcenters[i - CUDA_GRABCUT_K] / indexs[0][i];
- count += abs(indexs[0][i] - meansbg.length[i - CUDA_GRABCUT_K]);
- meansbg.length[i - CUDA_GRABCUT_K] = indexs[0][i];
- }
- delta = count;
- }
- }
- void updateCluster_gpu(PtrStepSz<uchar4> source, PtrStepSz<uchar> clusterIndex, PtrStepSz<uchar> mask, float4* kencter,
- int* kindexs, kmeansI& meansbg, kmeansI& meansfg, int& delta, cudaStream_t stream = nullptr)
- {
- dim3 grid(cv::divUp(source.cols, block.x), cv::divUp(source.rows, block.y));
- updateCluster<BLOCK_X, BLOCK_Y> <<<grid, block, 0, stream>>> (source, clusterIndex, mask, kencter, kindexs);
- updateCluster<BLOCK_X, BLOCK_Y> <<<1, block, 0, stream>>> (kencter, kindexs, meansbg, meansfg, delta, grid.x*grid.y);
- }
- void KmeansCuda::compute(GpuMat source, GpuMat clusterIndex, GpuMat mask, int threshold)
- {
- int delta = threshold + 1;
- initKmeans_gpu(*bg, *fg);
- while (delta> threshold)
- {
- findNearestCluster_gpu(source, clusterIndex, mask, *bg, *fg);
- updateCluster_gpu(source, clusterIndex, mask, dcenters, dlenght, *bg, *fg, *d_delta);
- // 可以每隔五次调用一次这个
- cudaMemcpy(&delta, d_delta, sizeof(int), cudaMemcpyDeviceToHost);
- }
- }
- updateCluster
这里肯定有点反直觉, 十多行代码扩展到百多行, 效率还增加了十倍, 但是 GPU 运算就是这样, 每步能多利用核心就行, 还有, cudaMemcpy 花费真的大, 这里有个失败的尝试, 因为在这基础上, 都是用更多的计算核心带来的好处, 就想把这个算法再扩展了下, 前景与背景二个各分五类, 一共十类, 故当前准备用 block 的 z 为 10, 当做每块的索引, 和我想象不同, 这个要的时间在 0.4ms 左右, 时间关系, 没有继续测试, 后续针对这块再仔细分析比对下.
如下显示的效果图.
然后就是 GMM 的实现, 这个实现和 k-means 过程比较相似, 代码就不贴了, 当时效果完成后, 我的主要疑问在如果确认我的实现是对的, 但是下面这张图出来后, 我就知道应该是对了. GMM 对 k-means 每块求的结果, 然后用这结果来得到颜色的分类, 这二者效果应该差不多是一样, 但是肯定又有点区别 (image 是 k-means 的结果, seg image 是 GMM 一次 learning 然后再分配的结果), 当然最后确认还是我比较 opencv 里 GMM 里各个分类与我实现 GMM 的值的比较, 一般来说, 各个分类的平均值与个数相差不大, 各个分类在所有点的比例类似, 当然二者的结果不会完全一样, opencv 中 k-means 中用 k-means++ 随机选择五个差异比较大的中心, 就算用 opencv 自己重复算同一张图, 每次结果都会有细小区别.
最后是 maxflow 的实现, 这个过程我参照了 https://web.iiit.ac.in/~vibhavvinet/Software/ 的实现, 要说这份代码稍稍改下还是能运行的, 就是数据来源全在一个文件上, 包含 formsource,tosink,edge 的信息, 主要查看 push 与 relabel 与 bfs 的实现, 可以去掉一些很多无关紧要的代码, 并且与 opencv 对接, 所有显存全用 gpumat 代码, 差不多只有 200 行左右的代码. 记的用这里的数据如 person/sponge 大约 50 多次迭代 push-relabel 后就能得到很好的结果 (50 次 640*480 类似的图只需要 5-8ms, 如果是 416*416 会更低, 大约 0.07ms 就能一次迭代), 在这给了我信心能继续做下去, 当然后续结果另说.
后续就是一个组合过程, 组合过程参照 opencv 的实现.
- void GrabCutCude::renderFrame(GpuMat x3source, cv::Rect rect)
- {
- cudaDeviceSynchronize();
- rgb2rgba_gpu(x3source, source);
- setMask_gpu(mask, rect);
- kmeans->compute(source, clusterIndex, mask, 0); //testShow(showSource, clusterIndex, "image");
- gmm->learning(source, clusterIndex, mask); //testShow(showSource, clusterIndex, "seg image");
- int index = 0;
- while (index++ <iterCount)
- {
- gmm->assign(source, clusterIndex, mask);
- gmm->learning(source, clusterIndex, mask);
- graph->addTermWeights(source, mask, *(gmm->bg), *(gmm->fg), lambda);
- graph->addEdges(source, gamma);
- graph->maxFlow(mask, maxCount);
- }
- }
- renderFrame
其中有几处要求聚合运算, 分别是求 maxflow 边的权重时, 整张图的 beta 值, k-means 更新聚合, gmm 的训练, 其中求 beta 只费了 0.02ms, 而 k-means 在 0.12ms, 而 gmm 的训练在 1.2ms 左右, beta 只是把所有数据相加, 每个像素对应 4byte 的共享存储, 而 k-means 把数据取成十个类别, 每个像素对应对应 4byte*(3+1)*10 个字节, 时间花费 0.02/0.12 应该是一个正常范围, 而 gmm 与 k-means 运算差不多完全一样, 只是共享显存占用 4byte*(3+3*3+1)*10 个字节, 时间花费比 k-means 多了十倍. 这个比较很有意思, 共享显存用越多, 可以看到, GPU 利用率会变低, 求 beta 时, occupancy 可以到达 100%, 而在 k-meansk 中只有 25%,updateGMM 只有 6.25%. 从前面 k-means 的代码看, 每个聚合分成二部分, 经调整 GMM 的前后部分发现, 上部分降低线程块 block 可以降低时间, 而下部分升高线程块 block 可以降低时间, 其实这么算的话 GMM 的计算可以考虑分开算, 让 k-means 的花费来看, 最多 0.3ms 左右, 1.2ms 肯定还有很多优化空间.
如下结合我们自己编译的 darknet 动态链接库组合的代码.
- void testYoloGrabCutCuda()
- {
- cv::VideoCapture cap;
- cv::VideoWriter video;
- cv::Mat frame;
- cv::Mat result, reframe;
- cap.open(0);
- char* modelConfiguration = "yolov3.cfg";
- char* modelWeights = "yolov3.weights";
- network *net = load_network(modelConfiguration, modelWeights, 0);
- layer l = net->layers.NET->n - 1];
- set_batch_network.NET, 1);
- cv::Mat netFrame.NET->h, net->w, CV_32FC3);
- const char* window_name = "yolo3";
- //namedWindow(window_name);
- cv::namedWindow("image");
- cv::namedWindow("seg image");
- createUI("1 image");
- int height = net->w;
- int width = net->h;
- GrabCutCude grabCut;
- grabCut.init(width, height);
- GpuMat gpuFrame;
- while (cv::waitKey(1) <1)
- {
- updateUI("1 image", grabCut);
- cap>> frame;
- cv::resize(frame, reframe, cv::Size.NET->w, net->h));
- //reframe.convertTo(netFrame, CV_32FC3, 1 / 255.0);
- image im = mat_to_image(reframe);
- float *predictions = network_predict.NET, im.data);
- int nboxes = 0;
- float thresh = 0.7f;
- detection curdet = {};
- float maxprob = thresh;
- bool bFind = false;
- detection *dets = get_network_boxes.NET, im.w, im.h, thresh, 0, 0, 0, &nboxes);
- for (int i = 0; i <nboxes; i++)
- {
- auto det = dets[i]; //0 person 39 bottle
- if (det.prob[39]> maxprob)
- {
- maxprob = det.prob[39];
- curdet = det;
- bFind = true;
- }
- }
- if (!bFind)
- continue;
- int width = curdet.bbox.w + 20;
- int height = curdet.bbox.h + 20;
- cv::Rect rectangle(curdet.bbox.x - width / 2.f, curdet.bbox.y - height / 2.f, width, height);
- if (bCpu)
- {
- cv::Mat bgModel, fgModel;
- cv::grabCut(reframe, result, rectangle, bgModel, fgModel, iterCount, cv::GC_INIT_WITH_RECT);
- }
- else
- {
- grabCut.setParams(iterCount, gamma, lambda, maxCount);
- cudaDeviceSynchronize();
- gpuFrame.upload(reframe);
- grabCut.renderFrame(gpuFrame, rectangle);
- grabCut.mask.download(result);
- }
- cv::compare(result, cv::GC_PR_FGD, result, cv::CMP_EQ);
- cv::Mat foreground(reframe.size(), CV_8UC3, cv::Scalar(0, 0, 0));
- reframe.copyTo(foreground, result);
- cv::imshow("image", foreground);
- cv::rectangle(reframe, rectangle, cv::Scalar(0, 0, 255), 3);
- cv::imshow("seg image", reframe);
- }
- }
- testYoloGrabCutCuda
这段代码就是前面显示的结果, 但是最后我发现, 效果并不好, 至少相对我参考那段 maxflow 效果来说, 迭代五十次只需要 5ms, 再稍微优化下, k-means 加上 GMM 可以控制在 3ms 内, 这样加上 yolo3 也可以控制在 30 桢左右, 但是, 就前面那个杯子来说, 需要迭代超过 300 次才能有个比较好的效果, 如果是人, 需要迭代上千次才能达到 opencv 的效果, 那就没啥意义了, 后续准备的优化与整合进引擎也暂停.
至于为啥会这样, 我分析了下, 原参考 push-relabel 里数据是用种子生成的方式来生成 formsource,tosink, 这样 formsource 与 tosink 本来就接近结果了.
而按照 opencv 里的流程, 出来的 formsource 与 tosink 图如下.
这样就需要迭代更多次的 push_relabel 来完成, 主要因为这二次方式生成的 GMM 有很大区别, 种子点生成的 GMM 各个分类与占比本来就好, 我试验了下, 扩大生成的边框, 马上 push_relabel 的迭代次数加个百次, 这是因为前景中混成来的背景越多, 背景分类与占比越大, 不确定区颜色通过 GMM 生成的 formsource 与 tosink 差值越小, 虽然 push-relabel 每个点可以独立计算, 很适合 GPU 算, 但是他的 push 与 relabel 方式导致他每次可能就少数点在流动, 大多点根本不满足要求.
那么要考虑在 100 次迭代内得到比较好的效果, 需要种子点方式, 加上一些特定需求, 比如背景与人变动不大, 只是人在背景中的位置在不断变化, 可以先用几桢算 mask-rcnn 得到 mask, 用 mask 得到相对准确的 GMM 模型, 再把这个 GMM 模型算对应 graph 的点权重, 然后用 maxflow 算的比较清晰的边, 最后用导向滤波还原大图, 这模式对现在的流程有些改动, 故此记录下现在流程, 如果后续有比较好的效果再来继续说.
2019/3/20 号更新:
为了验证我的想法以及评论区有同学提出 formsource 与 tosink 用同一例子比较, 深以为然. 下面是先用几桢效果的图算出 GMM, 然后用这个 GMM 值直接算 formsource 与 tosink, 如下是对应的 formsource 与 tosink 图.
- void GrabCutCude::renderFrame(GpuMat x3source)
- {
- rgb2rgba_gpu(x3source, source);
- graph->addEdges(source, gamma);
- int index = 0;
- while (index++ <iterCount)
- {
- graph->addTermWeights(source, *(gmm->bg), *(gmm->fg), lambda);
- graph->maxFlow(mask, maxCount);
- }
- }
- renderFrame
相对于原先方框里算 k-means,GMM 模型, 这个是在已经得到 GMM 模型的情况下, 用 GMM 直接算 formsource 与 tosink, 可以看到, 相对于原来来说, 结果已经很干净了, 这种情况下, 50 次 push-relabel 就能得到比较准确的结果.
来源: https://www.cnblogs.com/zhouxin/p/10567052.html