example-apps/PDD/pcb-defect-detection/libs/box_utils/cython_utils/nms.pyx

   1 # --------------------------------------------------------
   2 # Fast R-CNN
   3 # Copyright (c) 2015 Microsoft
   4 # Licensed under The MIT License [see LICENSE for details]
   5 # Written by Ross Girshick
   6 # --------------------------------------------------------
   7
   8 import numpy as np
   9 cimport numpy as np
  10
  11 cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
  12     return a if a >= b else b
  13
  14 cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
  15     return a if a <= b else b
  16
  17 def nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
  18     cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]
  19     cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]
  20     cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]
  21     cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3]
  22     cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]
  23
  24     cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  25     cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1]
  26
  27     cdef int ndets = dets.shape[0]
  28     cdef np.ndarray[np.int_t, ndim=1] suppressed = \
  29             np.zeros((ndets), dtype=np.int)
  30
  31     # nominal indices
  32     cdef int _i, _j
  33     # sorted indices
  34     cdef int i, j
  35     # temp variables for box i's (the box currently under consideration)
  36     cdef np.float32_t ix1, iy1, ix2, iy2, iarea
  37     # variables for computing overlap with box j (lower scoring box)
  38     cdef np.float32_t xx1, yy1, xx2, yy2
  39     cdef np.float32_t w, h
  40     cdef np.float32_t inter, ovr
  41
  42     keep = []
  43     for _i in range(ndets):
  44         i = order[_i]
  45         if suppressed[i] == 1:
  46             continue
  47         keep.append(i)
  48         ix1 = x1[i]
  49         iy1 = y1[i]
  50         ix2 = x2[i]
  51         iy2 = y2[i]
  52         iarea = areas[i]
  53         for _j in range(_i + 1, ndets):
  54             j = order[_j]
  55             if suppressed[j] == 1:
  56                 continue
  57             xx1 = max(ix1, x1[j])
  58             yy1 = max(iy1, y1[j])
  59             xx2 = min(ix2, x2[j])
  60             yy2 = min(iy2, y2[j])
  61             w = max(0.0, xx2 - xx1 + 1)
  62             h = max(0.0, yy2 - yy1 + 1)
  63             inter = w * h
  64             ovr = inter / (iarea + areas[j] - inter)
  65             if ovr >= thresh:
  66                 suppressed[j] = 1
  67
  68     return keep
  69
  70 def nms_new(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
  71     cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]
  72     cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]
  73     cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]
  74     cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3]
  75     cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]
  76
  77     cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  78     cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1]
  79
  80     cdef int ndets = dets.shape[0]
  81     cdef np.ndarray[np.int_t, ndim=1] suppressed = \
  82             np.zeros((ndets), dtype=np.int)
  83
  84     # nominal indices
  85     cdef int _i, _j
  86     # sorted indices
  87     cdef int i, j
  88     # temp variables for box i's (the box currently under consideration)
  89     cdef np.float32_t ix1, iy1, ix2, iy2, iarea
  90     # variables for computing overlap with box j (lower scoring box)
  91     cdef np.float32_t xx1, yy1, xx2, yy2
  92     cdef np.float32_t w, h
  93     cdef np.float32_t inter, ovr
  94
  95     keep = []
  96     for _i in range(ndets):
  97         i = order[_i]
  98         if suppressed[i] == 1:
  99             continue
 100         keep.append(i)
 101         ix1 = x1[i]
 102         iy1 = y1[i]
 103         ix2 = x2[i]
 104         iy2 = y2[i]
 105         iarea = areas[i]
 106         for _j in range(_i + 1, ndets):
 107             j = order[_j]
 108             if suppressed[j] == 1:
 109                 continue
 110             xx1 = max(ix1, x1[j])
 111             yy1 = max(iy1, y1[j])
 112             xx2 = min(ix2, x2[j])
 113             yy2 = min(iy2, y2[j])
 114             w = max(0.0, xx2 - xx1 + 1)
 115             h = max(0.0, yy2 - yy1 + 1)
 116             inter = w * h
 117             ovr = inter / (iarea + areas[j] - inter)
 118             ovr1 = inter / iarea
 119             ovr2 = inter / areas[j]
 120             if ovr >= thresh or ovr1 > 0.95 or ovr2 > 0.95:
 121                 suppressed[j] = 1
 122
 123     return keep