import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle

from itertools import cycle
cycol = cycle('bgrcmk')

# 資料準備
dets = np.random.rand(3, 5) + [0,0,1,1,0]
dets /= 2 # 這個是為了可以在圖裡畫出來。

def py_cpu_nms(dets,
 
 thresh):
    """Pure Python NMS baseline."""
    
    for i in range(dets.shape[0]):
        a,b,c,d,e = dets[i]
        plt.gca().add_patch(
                    plt.Rectangle((a,b),c - a,d - b, facecolor = 'green', fill = False,
                                  edgecolor='r', linewidth=3)
                    )
 

    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # n x 1
    print areas.shape, areas
    order = scores.argsort()[::-1] # n x 1
    print order

    keep = []
    while order.size > 0:
        i = order[
 
0]
        keep.append(i)
        
        # 以置信度最高的 x_bottom_left 為標準，找不小於它的
        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:]])
        
        for i in range(len(xx1)):
            plt.gca().add_patch(
                    plt.Rectangle((xx1[i],yy1[i]),xx2[i] - xx1[i],yy2[i]- yy1[i], facecolor = 'black', fill = False,
                                  edgecolor=cycol.next(), linewidth=3)
                    )
        # 計算置信度最大的矩形與其它矩形相交的面積
        w = np.maximum(0.0, xx2 - xx1 + 1) 
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        
        # 計算 Intersection over Union(IoU)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        
        # 獲得 從 order[1] 開始的所有滿足條件的矩形的下標
        inds = np.where(ovr <= thresh)[0]
        
        # 因為 order[0] 已經包含在 keep 裡面了，所以 inds 要做一個向右的 1 的偏移。
        order = order[inds + 1]

        
    plt.show()
    return keep


py_cpu_nms(dets, 1)

因為資料是隨機生成的，所以上面程式碼執行的結果可能和下圖不同。