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)