NMS(非极大值抑制)的python,cpu,gpu实现

必要性

NMS(非极大值抑制)是目标检测中用来确定最佳检测框的手段,根据目标检测流程,若果没有NMS步骤,其每个检测框都会有大量重叠度很高的预测框表示同一个目标。如下图:

NMS(非极大值抑制)的python,cpu,gpu实现_第1张图片NMS(非极大值抑制)的python,cpu,gpu实现_第2张图片

 左图为经过NMS的预测结果,右图为未经过NMS的结果,很明显,左图才是我们需要的结果。

过程

以yolo为例,其预测结果tensor为(bs,boxes,location(4)+confidence+num_classes)的形式。

NMS的python代码实现:

def non_max_supperression(boxes, num_classes, conf_thres=0.5, nms_thres=0.4):
    bs = np.shape(boxes)[0]

    #boxes中的location是(中心x,中心y, 宽w, 高h)形式的,需要调整为(左上x,左上y,右下x,右下y)的形式,方便IOU的计算。
    shape_boxes = np.zeros_like(boxes[:,:,:4])
    shape_boxes[:,:,0] = boxes[:,:,0] - boxes[:,:,2] / 2
    shape_boxes[:,:,1] = boxes[:,:,0] - boxes[:,:,3] / 2
    shape_boxes[:,:,2] = boxes[:,:,0] + boxes[:,:,2] / 2
    shape_boxes[:,:,3] = boxes[:,:,0] + boxes[:,:,3] / 2
    #替换
    boxes[:,:,:4] = shape_boxes;
    
    output=[]
    #对每张图进行处理
    for i in range(bs):
        #prediction与boxes相比,bs维度没有了:(num_box,4+1+num_class)
        prediction = boxes[i];
        #取出置信度
        score = prediction[:,4]
        #如果置信度低于阈值,则直接忽略,高于阈值的才进行下一步的抑制
        mask = score > conf_thres
        detction = prediction[mask]
        #取预测框的类别和对应的预测概率
        class_conf = np.expand_dims(np.max(detction[:,5:],axis=-1),axis=-1)
        class_pred = np.expand_dims(np.argmax(detection[:,5:],axis=-1),axis=-1)
        #predction进行重组,此时的prediction为[num_box,4(转换后的框位置信息)+1(有目标的置信        
        #  度)+2(目标的类别置信度和类别归属)
        prediction = np.concatenate([prediction[:,:5],class_conf,class_pred],-1)
        #图中都有哪些类别
        unique_class = np.unique(detection[:,-1])
        if len(unique_class) == 0:
            continue;
        
        #列表对抑制后的box进行存储
        best_box = []
        
        for c in unique_class:
            cls_mask = detection[:,-1] == c
            #选出了该类别的detection
            detection_chosen = detection[cls_mask]
            #根据是否有目标的置信度进行排序
            scores = detction_chosen[:,4]
            arg_sort = np.argsort(scores)[::-1]//np.argsort升序排,取反进行降序排
            detction_chosen = detction_chosen[arg_sort]

            while len(detction_chosen) != 0:
                #将具有最大置信度的框加入结果框列表中
                best_box.append(detction_chosen[0])
                #框只有一个,则不用筛选了
                if len(detction_chosen)==1:
                    break
                #计算该框与其他同类框的交并比iou
                ious = iou(best_box[-1],detction_chosen[1:])
                #筛选剩下交并比小于设定阈值的检测框,因为交并比小的框可能来自其他同类目标,而交                
                  并比过大的应当是对同一个目标的重复预测
                detction_chosen = detction_chosen[1:][ious

部署时需后处理NMS通过c在CPU或GPU上完成,首先是CPU:

//Box为自定义类,其成员变量有left,top,right,bottom,confidence和label
vector cpu_decode(float* predict, int rows, int cols, float con_thres=0.25f, float         
                       nms_thres = 0.45f){
    //predict为python训练端保存过来的预测结果向量,boxes用来存储坐标转换后的中间向量
    vector boxes;
    //向量前五列为4(框信息)+1(是否有目标的置信度),不包含类别信息
    int num_classes = cols - 5;
    //对每个框进行处理
    for(int i=0; ib.confidence;});
    //定义remove_flags标记boxes中哪些框是需要保留的,哪些框是要删除的
    vector remove_flags(boxes.size());
    //用来存储抑制后的结果框
    vector box_result;
    //预分配内存,防止不断动态分配内存带来的耗时
    box_result.reserve(boxes.size());
    
    //iou计算函数
    auto iou = [](const Box& a,const Box& b){
        //获取相交矩形的左上、右下坐标
        float cross_left = std::max(a.left,b.left);
        float cross_top = std::max(a.top,b.top);
        float cross_right = std::min(a.right,b.right);
        float cross_bottom = std::min(a.bottom,b.bottom);
        //计算相交矩形的面积
        float cross_area = std::max(0.0f, cross_right-cross_left) 
                           * std::max(0.0f, cross_bottom-cros_top));
        //计算相并面积
        float union_area = std::max(a.right-a.left,0.0f) * std::max(a.bottom-a.top,0.0f)
                         + std::max(b.right-b.left,0.0f) * std::max(b.bottom-b.top,0.0f)
                         - cross_area;
        if(cross_area || union_area == 0) return 0.0f;
        
        return cross_area / union_area;
    };
    
    for(int i=0;i= nms_thres){
                    remove_flags[j] = true;    
                }
            }
        }
    }
    return box_result;
}

CUDA编程到GPU端,在CPU上做一定的修改:

//参数域cpu含义相同
vector gpu_decode(float* predict, int rows, int cols, float con_thres=0.25f, float         
                       nms_thres = 0.45f){
    
    vector box_result;
    //创建流,一般在开头就创建,这里为了演示说明,在这里创建
    cudaStream_t stream=nullptr;
    checkRuntime(cudaStreamCreate(&stream));//checkRuntime为略作修改的检查代码,检查正常创建
    
    //显卡上的传入预测框
    float* predict_device = nullptr;
    //显卡上处理后的结果
    float* output_device = nullptr;
    //显卡上处理后的结果传到host上
    float* output_host = nullptr;

    int max_objects = 1000;
    int NUM_BOX_ELEMENT = 7;//left,top,right,bottom,confidence,class,keepflag(是否去除的flag
    
    //分配global memory用来接收从host传过来的predict
    checkRuntime(cudaMalloc(&predict_device, rows*cols*sizeof(float));
    //分配global memory用来存储处理后的目标信息
    checkRuntime(cudaMalloc(&output_device, max_objects * NUM_BOX_ELEMENT + sizeof(float));
    //分配pinned memory用来从设备到host
    checkRuntime(cudaMallocHost(&output_host, max_objects*NUM_BOX_ELEMENT+sizeof(float));
    
    //异步复制
    checkRuntime(cudaMemAsync(predict_device,predict,rows*cols**sizeof(float),
                              cudaMemcpyHostToDevice,stream)
    //框解码和nms核函数的启动函数
    decode_kernel_invoker(
        predict_device,rows,cols-5,conf_thres,nms_thres,nullptr,output_device,
        max_objects,NUM_BOX_ELEMENT,stream
    );
    
    //将gpu上的预测框拷贝到host
    checkRuntime(cudaMemcpyAsync(output_host,output_device,
                                 sizeof(int)+max_objects*NUM_BOX_ELEMENT*sizeof(float),
                                 cudaMemcpyDeviceToHost,stream
    ));
    checkRuntime(cudaStreamSynchronize(stream));
    
    int num_boxes = min((int)output_host[0],max_objects);
    //将结果框存入box_result
    for(int i=0;i512 ? 512:num_bboxes;//block一般取1024下较大的32倍数
    //相当于向上取整,要达到的目的是grid*block大于等于num_bboxes且被整除
    auto grid = (num_bboxes + block - 1) / block;
    
    //调用框解码核函数
    decode_kernel<<>>(
        predict, num_bboxes, num_classes, conf_thres,invert_affine_matrix,parray,
        max_objects,NUM_BOX_ELEMENT
    );

    //确定线程的配置参数,开启线程数为max_objects个数
    auto block = num_bboxes>512 ? 512:max_objects;//block一般取1024下较大的32倍数
    //相当于向上取整,要达到的目的是grid*block大于等于num_bboxes且被整除
    auto grid = (max_objects+ block - 1) / block;
    //进行多线程nms
    fast_nms_kernel<<>>        
                   (parray,max_objects,nms_thres,NMU_BOX_ELEMENT);
    //parray中的count可能会超出max_objects,因为线程所有线程会一直执行到那一步,虽然会不符合条    
    件从而不往下走,但count会一直累加,因此后面要取最小值
}

static __global__ void decode_kernel(
    float* predict,int num_bboxes,int num_classes,float conf_thres,
    float* invert_affine_matrix,float* parray,int max_objects,int NUM_BOXELEMENT
){
    int position = blockDim.x * blockIdx.x + threadIdx.x;
    if (position >= num_bboxes) return;
    
    //获取每个框(对应一个线程)的首地址
    float* pitem = predict+(5+num_classes)*position;
    //获取有无目标的置信度
    float objectness = pitem[4];
    if(objectness < conf_thres) return;
    //获取表示类别信息的地址
    float* class_confidence = pitem+5;
    //获取当前对于当前指向类别的置信度,并指向下一个类别
    float confidence = *class_confidence++;
    int label=0;
    //其实是在找类别中最大置信度的作为预测的类别
    for(int i=1;i confidence){
            confidence = *class_confidence;
            label = i;
        }    
    }
    
    confidence* = objectness;
    if(confidence < conf_thres) return;
    
    //能执行到这说明得到了置信度足够的目标框,需要对这个框进行解码,并将解码信息存入parray
    //parray = count,box1,box2,box3
    //atomicAdd -> count+=1 返回的是old_count,新的值被存入内存中
    int index = atomicAdd(parray,1);//
    if(index >= max_objects) return;
    
    //获取坐标信息并转化
    float cx =     *pitem++;
    float cy =     *pitem++;
    float width =  *pitem++;
    float height = *pitem++;
    float left =   cx-0.5f*width;
    float top =    cy-0.5f*height;
    float right =  cx+0.5f*width;
    float bottom = cy+0.5f*height;
    
    //将转换后的left,top,right,bottom,confidence,class,keepflag填入parray
    float* pout_item = parray + 1 + index * NUM_BOX_ELEMENT;
    *pout_item++ = left;
    *pout_item++ = top;
    *pout_item++ = right;
    *pout_item++ = bottom;
    *pout_item++ = confidence;
    *pout_item++ = label;
    *pout_item++ = 1;//keepflag为1时表示保持,不删除
}

//测试mAP用cpuNMS
//日常推理可用GPU
//GPU上的NMS其实相当于开了框数量个线程,每个线程循环了框数量次(进行比较)
static __global__ void fast_nms_kernel(float* bboxes,int max_objects,float thres,
                                       int NUM_BOX_ELEMENT){
    int position = blockDim.x*blockIdx.x+threadIdx.x;
    //即decode中的index,剩下的框个数,多余的线程不需要工作
    int count = min((int)*bboxes,max_objects);
    if(position >= count) return;    
    //获取当前框,对flag的操作一定在pcurrent上完成,因为pcurrent对应的是当前线程,而不是pitem   
    float* pcurrent = bboxes + 1 + position * NUM_BOX_ELEMENT;
    //去除条件:重叠度高、类别相同、置信度低小于已有的
    for(int i=0;i= pcurrent[4]){
            //其他框置信度与当前框相同并且其为当前框之前的框,那么当前框保留,因为之前的框在其他            
              线程中与当前框的比较中被执行到下一步进行了筛选,可能被打上了移除的印记,所有线程统    
              一为默认保留后面的框,这样要删除的话前面的框已经被打上删除记号了
            if(pitem[4]==pcurrent[4] && i thres){
                pcurrent[6] = 0;
                return;
            }
        }
    }
}
//只能在gpu中调用设备函数
static __device__ float box_iou(
        float aleft,float atop,float aright,float abottom,
        float bleft,float btop,float bright,float bbottom){
    //获取相交矩形的坐标
    float cleft = max(aleft,bleft);
    float ctop = max(atop,btop);
    float cright = min(aright,bright);
    float cbottom = max(abottom,bbottom);
    //相交矩形的面积
    float c_area = max(cright-cleft,0.0f) * max(cbottom-ctop,0.0f);
    if(c_area==0.0f) return 0.0f;

    //并
    float a_area = max(0.0f, aright-a_left)*max(0.0f,abottom-atop)
    float b_area = max(0.0f, bright-b_left)*max(0.0f,bbottom-btop)
    return c_area/(a_area+b_area-c_area);
}

以上核函数入口、核函数和设备函数需要nvcc编译,单独写cu文件不在cpp中。

你可能感兴趣的:(CUDA,目标检测,深度学习,计算机视觉)