Tensorflow版Faster RCNN源码解析（TFFRCNN） （05） nms_wrapper.py

本blog为github上CharlesShang/TFFRCNN版源码解析系列代码笔记

---------------我的学习笔记---------------

----------------本文做者吴疆--------------

------点击此处连接至博客园原文------

 

1.def nms(dets，thresh，force_cpu=False)

def nms(dets, thresh, force_cpu=False): """Dispatch to either CPU or GPU NMS implementations."""
    if dets.shape[0] == 0: return [] # 默认USE_GPU_NMS = True
    if cfg.USE_GPU_NMS and not force_cpu: return gpu_nms(dets, thresh, device_id=cfg.GPU_ID)  # gpu_nms.so
    else: return cpu_nms(dets, thresh)

选择以GPU或CPU模式执行nms，实际是.so动态连接对象执行，其中，dets是某类box和scor构成的数组，shape为（None，5），被nms_wrapper(...)函数调用

2.def nms_wrapper(scores，boxes，threshold = 0.7，class_sets = None)

# train_model()调用时未传入class_sets参数
def nms_wrapper(scores, boxes, threshold = 0.7, class_sets = None):  # box得分必须大于0.7
    # scores：R * num_class
    # boxes: R * (4 * num_class)
    # return: a list of K-1 dicts, no background, each is {'class': classname, 'dets': None | [[x1,y1,x2,y2,score],...]} 
    num_class = scores.shape[1] if class_sets is None else len(class_sets) assert num_class * 4 == boxes.shape[1],\ 'Detection scores and boxes dont match' class_sets = ['class_' + str(i) for i in range(0, num_class)] if class_sets is None else class_sets # class_sets = [class_0, class_1, class_2...]
    res = [] # 针对各种，构造该类的dets（含box坐标和score）
    for ind, cls in enumerate(class_sets[1:]): ind += 1  # skip background
        cls_boxes = boxes[:, 4*ind: 4*(ind+1)] cls_scores = scores[:, ind] dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32) keep = nms(dets, thresh=0.3)  # nms阈值为0.3
        dets = dets[keep, :]  # 类内nms处理
        dets = dets[np.where(dets[:, 4] > threshold)]  # score必须超过阈值0.7
        r = {} if dets.shape[0] > 0: r['class'], r['dets'] = cls, dets else: r['class'], r['dets'] = cls, None res.append(r) # res为列表，每一个元素为某类别标号和dets构成的字典
    return res

最后返回的res为列表，列表中每一个元素为字典，每一个字典含某类标号如(class_1,class_2...)和dets(该类box的坐标和对应score，针对各种类内先nms后取得分超过阈值的box)