| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970 |
- 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)
|
