You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
71 lines
2.5 KiB
Cython
71 lines
2.5 KiB
Cython
10 months ago
|
|
||
|
|
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)
|