def nms(result, thresh): # dets:(m,5) thresh:scaler ls=[] noempty_nums=0 for r in result: l=len(r) ls.append(l) if l>0: noempty_nums+=0 if noempty_nums==1: return result #ls=[set(range(r)) for r in result] dets=np.vstack(result) # if len(result[1])==0: # dets=result[0] # l_bottle=len(result[0]) # l_other=0 # else: # dets=np.vstack([result[0],result[1]]) # l_bottle=len(result[0]) # l_other=[len(result[0]),len(dets)] x1 = dets[:, 0] y1 = dets[:, 1] x2 = dets[:, 2] y2 = dets[:, 3] areas = (y2 - y1 + 1) * (x2 - x1 + 1) scores = dets[:, 4] keep = [] index = scores.argsort()[::-1] while index.size > 0: i = index[0] # every time the first is the biggst, and add it directly keep.append(i) x11 = np.maximum(x1[i], x1[index[1:]]) # calculate the points of overlap y11 = np.maximum(y1[i], y1[index[1:]]) x22 = np.minimum(x2[i], x2[index[1:]]) y22 = np.minimum(y2[i], y2[index[1:]]) w = np.maximum(0, x22 - x11 + 1) # the weights of overlap h = np.maximum(0, y22 - y11 + 1) # the height of overlap overlaps = w * h ious = overlaps / (areas[i] + areas[index[1:]] - overlaps) idx = np.where(ious <= thresh)[0] index = index[idx + 1] # because index start from 1 start=0 det_classes=[] keep=np.array(keep) for l in ls: indexs=keep[np.where((start<=keep)&(keep<l+start))] #indexs=keep[start<=keep<l+start] det_classes.append(dets[indexs]) start=l+start # for index in keep: # for i,s in enumerate(ls): # if index in s: # print() # index_bottle=[index for index in keep if index<l_bottle] # index_=[index for index in keep if index>=l_bottle] # bottle=dets[index_bottle,:] # impurity=dets[index_other,:] # plastic_can # metal_can # result=[bottle,impurity,plastic_can,metal_can] return det_classes
类间的nms
最新推荐文章于 2024-07-04 17:49:05 发布