|
|
| import numpy as np |
| import cv2 |
| cimport numpy as np |
| cimport cython |
| cimport libcpp |
| cimport libcpp.pair |
| cimport libcpp.queue |
| from libcpp.pair cimport * |
| from libcpp.queue cimport * |
|
|
| @cython.boundscheck(False) |
| @cython.wraparound(False) |
| cdef np.ndarray[np.int32_t, ndim=2] _pse(np.ndarray[np.uint8_t, ndim=3] kernels, |
| np.ndarray[np.int32_t, ndim=2] label, |
| int kernel_num, |
| int label_num, |
| float min_area=0): |
| cdef np.ndarray[np.int32_t, ndim=2] pred |
| pred = np.zeros((label.shape[0], label.shape[1]), dtype=np.int32) |
|
|
| for label_idx in range(1, label_num): |
| if np.sum(label == label_idx) < min_area: |
| label[label == label_idx] = 0 |
|
|
| cdef libcpp.queue.queue[libcpp.pair.pair[np.int16_t,np.int16_t]] que = \ |
| queue[libcpp.pair.pair[np.int16_t,np.int16_t]]() |
| cdef libcpp.queue.queue[libcpp.pair.pair[np.int16_t,np.int16_t]] nxt_que = \ |
| queue[libcpp.pair.pair[np.int16_t,np.int16_t]]() |
| cdef np.int16_t* dx = [-1, 1, 0, 0] |
| cdef np.int16_t* dy = [0, 0, -1, 1] |
| cdef np.int16_t tmpx, tmpy |
|
|
| points = np.array(np.where(label > 0)).transpose((1, 0)) |
| for point_idx in range(points.shape[0]): |
| tmpx, tmpy = points[point_idx, 0], points[point_idx, 1] |
| que.push(pair[np.int16_t,np.int16_t](tmpx, tmpy)) |
| pred[tmpx, tmpy] = label[tmpx, tmpy] |
|
|
| cdef libcpp.pair.pair[np.int16_t,np.int16_t] cur |
| cdef int cur_label |
| for kernel_idx in range(kernel_num - 1, -1, -1): |
| while not que.empty(): |
| cur = que.front() |
| que.pop() |
| cur_label = pred[cur.first, cur.second] |
|
|
| is_edge = True |
| for j in range(4): |
| tmpx = cur.first + dx[j] |
| tmpy = cur.second + dy[j] |
| if tmpx < 0 or tmpx >= label.shape[0] or tmpy < 0 or tmpy >= label.shape[1]: |
| continue |
| if kernels[kernel_idx, tmpx, tmpy] == 0 or pred[tmpx, tmpy] > 0: |
| continue |
|
|
| que.push(pair[np.int16_t,np.int16_t](tmpx, tmpy)) |
| pred[tmpx, tmpy] = cur_label |
| is_edge = False |
| if is_edge: |
| nxt_que.push(cur) |
|
|
| que, nxt_que = nxt_que, que |
|
|
| return pred |
|
|
| def pse(kernels, min_area): |
| kernel_num = kernels.shape[0] |
| label_num, label = cv2.connectedComponents(kernels[-1], connectivity=4) |
| return _pse(kernels[:-1], label, kernel_num, label_num, min_area) |
|
|
