161 lines
5.4 KiB
Python
161 lines
5.4 KiB
Python
import numpy as np
|
|
from numpy.lib.stride_tricks import as_strided
|
|
|
|
def _realnms(dets, only_max=False, iou_thres=0.35):
|
|
|
|
"""
|
|
non-maximum suppression: if only_max, will ignore iou_thres and return largest score bbox.
|
|
dets: list[list[x, y, w, h]]
|
|
only_max: bool
|
|
iou_thres: float between (0,1)
|
|
"""
|
|
dets = np.array(dets)
|
|
|
|
if len(dets) == 0:
|
|
return []
|
|
scores = dets[:, 4]
|
|
order = np.argsort(scores)[::-1]
|
|
dets = dets[order, :]
|
|
if only_max:
|
|
return np.array([dets[0]])
|
|
|
|
x1, y1, w, h, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
|
|
x2, y2 = x1 + w - 1, y1 + h - 1
|
|
|
|
areas = w * h
|
|
|
|
order = scores.argsort()[::-1]
|
|
|
|
keep_real = []
|
|
tol = 0.1
|
|
|
|
while order.size > 0:
|
|
i = order[0]
|
|
keep_real.append(i)
|
|
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:]])
|
|
|
|
inter_w = np.maximum(0.0, xx2 - xx1 + 1)
|
|
inter_h = np.maximum(0.0, yy2 - yy1 + 1)
|
|
|
|
inter_area = inter_w * inter_h
|
|
iou = inter_area / (areas[i] + areas[order[1:]] - inter_area)
|
|
|
|
inds = np.where(iou <= iou_thres)[0]
|
|
order = order[inds + 1]
|
|
|
|
return dets[keep_real,:]
|
|
|
|
|
|
def _pool2d(A, kernel_size, stride, padding, pool_mode='max'):
|
|
'''
|
|
2D Pooling
|
|
|
|
Parameters:
|
|
A: input 2D array
|
|
kernel_size: int, the size of the window
|
|
stride: int, the stride of the window
|
|
padding: int, implicit zero paddings on both sides of the input
|
|
pool_mode: string, 'max' or 'avg'
|
|
'''
|
|
# Padding
|
|
A = np.pad(A, padding, mode='constant')
|
|
|
|
# Window view of A
|
|
output_shape = ((A.shape[0] - kernel_size)//stride + 1,
|
|
(A.shape[1] - kernel_size)//stride + 1)
|
|
kernel_size = (kernel_size, kernel_size)
|
|
A_w = as_strided(A, shape = output_shape + kernel_size,
|
|
strides = (stride*A.strides[0],
|
|
stride*A.strides[1]) + A.strides)
|
|
A_w = A_w.reshape(-1, *kernel_size)
|
|
|
|
# Return the result of pooling
|
|
if pool_mode == 'max':
|
|
return A_w.max(axis=(1,2)).reshape(output_shape)
|
|
elif pool_mode == 'avg':
|
|
return A_w.mean(axis=(1,2)).reshape(output_shape)
|
|
|
|
def _nms(heat,hmax=None, kernel=3):
|
|
# hmax = MaxPool2D(kernel, strides=1,padding='same')(heat)
|
|
# heat = tf.where(tf.equal(hmax, heat), heat, tf.zeros_like(heat))
|
|
if hmax is None:
|
|
hmax = np.zeros_like(heat)
|
|
b, h, w, c = heat.shape
|
|
for batch in range(b):
|
|
for channel in range(c):
|
|
hmax[batch,:,:,channel] = _pool2d(heat[batch,:,:,channel], kernel_size=kernel, stride=1, padding=1)
|
|
assert (heat.shape == hmax.shape)
|
|
# make non local max item zero
|
|
keep = heat==hmax
|
|
# heat[heat!=hmax] = 0
|
|
return heat*keep
|
|
|
|
def postprocess_(outputs, max_objects=100, score_thres=0.5,
|
|
scale=None, input_shape=None, w_ori=None, h_ori=None,
|
|
nms=True, iou_thres=0.35, mapping_func='linear', e2e_coco=False, **kwargs):
|
|
|
|
assert len(outputs) % 3 == 0
|
|
n_stage = len(outputs)//3
|
|
dets = []
|
|
batch_index = 0
|
|
for stage in range(n_stage):
|
|
|
|
"""get output tensor and anchor tensor"""
|
|
reg, cls, cts = outputs[stage], outputs[stage+n_stage], outputs[stage+n_stage*2]
|
|
|
|
"""get dimension info"""
|
|
b, nrows, ncols, nclasses = cls.shape
|
|
assert b==1
|
|
# calculate here or pass by parameter
|
|
stride = 2**int(np.log2(1.0*input_shape[0]/nrows)+0.5)
|
|
|
|
"""iterate over all anchors and select those with valid score"""
|
|
# batch is always 0 here
|
|
for i in range(nrows):
|
|
for j in range(ncols):
|
|
"""class_id is the indice of labels. 0 is not background"""
|
|
class_id = np.argmax(cls[batch_index, i, j, :])
|
|
score = np.sqrt(cls[batch_index, i, j, class_id] * cts[batch_index, i, j, 0])
|
|
|
|
if score > score_thres:
|
|
if mapping_func == 'exp':
|
|
l, t, r, b = np.exp(reg[0, i, j])
|
|
elif mapping_func == 'linear':
|
|
reg_relu = np.clip(reg[0, i, j], 0, 1e8)
|
|
l, t, r, b = (2**(3+stage)) * (reg_relu**2)
|
|
else:
|
|
assert 0
|
|
cx, cy = j*stride + stride//2, i*stride + stride//2
|
|
xmin, ymin, xmax, ymax = cx-l, cy-t, cx+r, cy+b
|
|
if e2e_coco:
|
|
class_id += 1
|
|
dets.append([xmin, ymin, xmax, ymax, score, class_id])
|
|
|
|
dets = np.asarray(dets)
|
|
if scale is not None:
|
|
dets[..., :4] = dets[..., :4]*scale
|
|
|
|
if w_ori is not None and h_ori is not None and np.size(dets)>0:
|
|
# clip bbox make it inside image
|
|
dets[..., :4] = np.clip(dets[..., :4], [0., 0., 0., 0.],
|
|
np.c_[w_ori, h_ori, w_ori, h_ori])
|
|
|
|
|
|
if len(dets) > 0:
|
|
dets = np.asarray(dets)
|
|
dets[..., 2:4] = dets[..., 2:4] - dets[..., :2]
|
|
if nms:
|
|
dets_real = _realnms(dets, only_max=max_objects == 1, iou_thres=iou_thres)
|
|
else:
|
|
dets_real = dets
|
|
else:
|
|
dets_real= []
|
|
# [[x1,y1,w,h],[x1,y1,w,h]]
|
|
dets_real = np.asarray(dets_real)
|
|
dets_real = dets_real[..., :]
|
|
return dets_real.tolist()
|
|
|