深度學習目标檢測常用工具型代碼：對檢測出來的結果單獨進行nms操作

p.s. 這裡的例子我說的都是航空影像目标，是以大家考慮的時候都要考慮成俯視圖的狀态。

多目标檢測的時候，有容易發生擁擠的類别，比如小汽車，也有不易擁擠的類别，比如籃球場。是以需要不同的nms門檻值，而檢測一次有時候消耗的時間太多，是以我的解決辦法是，檢測的時候nms=1都輸出出來，然後在通過單獨對檢測的txt結果進行nms操作，比如所有檔案nms0.3、nms0.4各種nms全都可以重新輸出來，這樣會比檢測的同時進行nms友善快很多。

# -*- coding: utf-8 -*-
"""
Created on Mon Apr  9 21:09:09 2018
對每一類檢測生成的prebbox.txt進行nms操作
1.讀入文本，讀入bbox
@author: ygx
"""
import os
import numpy as np

#nms的門檻值設定
nms_thresh = 0.5
#nms之前的路徑
srcpath = r'/media/ygx/B4FE-5315/足球場/nms之前'
#nms之後的儲存路徑
dstpath = r'/media/ygx/B4FE-5315/足球場/nms/nms0.5'

def custombasename(fullname):
    return os.path.basename(os.path.splitext(fullname)[0])

def GetFileFromThisRootDir(dir,ext = None):
  allfiles = []
  needExtFilter = (ext != None)
  for root,dirs,files in os.walk(dir):
    for filespath in files:
      filepath = os.path.join(root, filespath)
      extension = os.path.splitext(filepath)[1][1:]
      if needExtFilter and extension in ext:
        allfiles.append(filepath)
      elif not needExtFilter:
        allfiles.append(filepath)
  return allfiles

def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    #print('dets:', dets)
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    ## index for dets
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        #print (areas[i] + areas[order[1:]] - inter)
        order = order[inds + 1]

    return keep


def nmsbynamedict(nameboxdict, nms,thresh):
    nameboxnmsdict = {x: [] for x in nameboxdict}
    for imgname in nameboxdict:
        #print('imgname:', imgname)
        #keep = py_cpu_nms(np.array(nameboxdict[imgname]), thresh)
        #print('type nameboxdict:', type(nameboxnmsdict))
        #print('type imgname:', type(imgname))
        #print('type nms:', type(nms))
        keep = nms(np.array(nameboxdict[imgname]),thresh)
        #print('keep:', keep)
        outdets = []
        #print('nameboxdict[imgname]: ', nameboxnmsdict[imgname])
        for index in keep:
            # print('index:', index)
            outdets.append(nameboxdict[imgname][index])
        nameboxnmsdict[imgname] = outdets
    return nameboxnmsdict



def nmsbase(srcpath, dstpath, nms):
    filelist = GetFileFromThisRootDir(srcpath)
    for fullname in filelist:
        name = custombasename(fullname)
        #print('name:', name)
        dstname = os.path.join(dstpath, name + '.txt')
        with open(fullname, 'r') as f_in:
            nameboxdict = {}
            lines = f_in.readlines()
            splitlines = [x.strip().split(' ') for x in lines]
            for splitline in splitlines:
                 oriname = splitline[0]#擷取檔案名P0051
                 confidence = splitline[1]
                 poly = list(map(float, splitline[2:]))
                 det = poly             
                 det.append(confidence)
                 det = list(map(float, det))  
                 #det.append(oriname)#變到清單裡面 一一對應
                 if (oriname not in nameboxdict):
                     nameboxdict[oriname] = []
                 nameboxdict[oriname].append(det)
            nameboxnmsdict = nmsbynamedict(nameboxdict, nms ,nms_thresh)
            with open(dstname, 'w') as f_out:
                for imgname in nameboxnmsdict:
                    for det in nameboxnmsdict[imgname]:
                        #print('det:', det)
                        confidence = det[-1]
                        bbox = det[0:-1]
                        outline = str(imgname) + ' ' +str(confidence) + ' ' + ' '.join(map(str, bbox))
                        #outline = imgname + ' ' + str(confidence) + ' ' + ' '.join(map(str, bbox))
                        #print('outline:', outline)
                        f_out.write(outline + '\n')

if __name__ == '__main__':
    nmsbase(srcpath,
              dstpath,
              py_cpu_nms)