101 lines
4.1 KiB
Python
101 lines
4.1 KiB
Python
import numpy as np
|
|
from config import INPUT_SIZE
|
|
|
|
_default_anchors_setting = (
|
|
dict(layer='p3', stride=32, size=48, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
|
dict(layer='p4', stride=64, size=96, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
|
dict(layer='p5', stride=128, size=192, scale=[1, 2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
|
)
|
|
|
|
|
|
def generate_default_anchor_maps(anchors_setting=None, input_shape=INPUT_SIZE):
|
|
"""
|
|
generate default anchor
|
|
|
|
:param anchors_setting: all informations of anchors
|
|
:param input_shape: shape of input images, e.g. (h, w)
|
|
:return: center_anchors: # anchors * 4 (oy, ox, h, w)
|
|
edge_anchors: # anchors * 4 (y0, x0, y1, x1)
|
|
anchor_area: # anchors * 1 (area)
|
|
"""
|
|
if anchors_setting is None:
|
|
anchors_setting = _default_anchors_setting
|
|
|
|
center_anchors = np.zeros((0, 4), dtype=np.float32)
|
|
edge_anchors = np.zeros((0, 4), dtype=np.float32)
|
|
anchor_areas = np.zeros((0,), dtype=np.float32)
|
|
input_shape = np.array(input_shape, dtype=int)
|
|
|
|
for anchor_info in anchors_setting:
|
|
|
|
stride = anchor_info['stride']
|
|
size = anchor_info['size']
|
|
scales = anchor_info['scale']
|
|
aspect_ratios = anchor_info['aspect_ratio']
|
|
|
|
output_map_shape = np.ceil(input_shape.astype(np.float32) / stride)
|
|
output_map_shape = output_map_shape.astype(np.int)
|
|
output_shape = tuple(output_map_shape) + (4,)
|
|
ostart = stride / 2.
|
|
oy = np.arange(ostart, ostart + stride * output_shape[0], stride)
|
|
oy = oy.reshape(output_shape[0], 1)
|
|
ox = np.arange(ostart, ostart + stride * output_shape[1], stride)
|
|
ox = ox.reshape(1, output_shape[1])
|
|
center_anchor_map_template = np.zeros(output_shape, dtype=np.float32)
|
|
center_anchor_map_template[:, :, 0] = oy
|
|
center_anchor_map_template[:, :, 1] = ox
|
|
for scale in scales:
|
|
for aspect_ratio in aspect_ratios:
|
|
center_anchor_map = center_anchor_map_template.copy()
|
|
center_anchor_map[:, :, 2] = size * scale / float(aspect_ratio) ** 0.5
|
|
center_anchor_map[:, :, 3] = size * scale * float(aspect_ratio) ** 0.5
|
|
|
|
edge_anchor_map = np.concatenate((center_anchor_map[..., :2] - center_anchor_map[..., 2:4] / 2.,
|
|
center_anchor_map[..., :2] + center_anchor_map[..., 2:4] / 2.),
|
|
axis=-1)
|
|
anchor_area_map = center_anchor_map[..., 2] * center_anchor_map[..., 3]
|
|
center_anchors = np.concatenate((center_anchors, center_anchor_map.reshape(-1, 4)))
|
|
edge_anchors = np.concatenate((edge_anchors, edge_anchor_map.reshape(-1, 4)))
|
|
anchor_areas = np.concatenate((anchor_areas, anchor_area_map.reshape(-1)))
|
|
|
|
return center_anchors, edge_anchors, anchor_areas
|
|
|
|
|
|
def hard_nms(cdds, topn=10, iou_thresh=0.25):
|
|
if not (type(cdds).__module__ == 'numpy' and len(cdds.shape) == 2 and cdds.shape[1] >= 5):
|
|
raise TypeError('edge_box_map should be N * 5+ ndarray')
|
|
|
|
cdds = cdds.copy()
|
|
indices = np.argsort(cdds[:, 0])
|
|
cdds = cdds[indices]
|
|
cdd_results = []
|
|
|
|
res = cdds
|
|
|
|
while res.any():
|
|
cdd = res[-1]
|
|
cdd_results.append(cdd)
|
|
if len(cdd_results) == topn:
|
|
return np.array(cdd_results)
|
|
res = res[:-1]
|
|
|
|
start_max = np.maximum(res[:, 1:3], cdd[1:3])
|
|
end_min = np.minimum(res[:, 3:5], cdd[3:5])
|
|
lengths = end_min - start_max
|
|
intersec_map = lengths[:, 0] * lengths[:, 1]
|
|
intersec_map[np.logical_or(lengths[:, 0] < 0, lengths[:, 1] < 0)] = 0
|
|
iou_map_cur = intersec_map / ((res[:, 3] - res[:, 1]) * (res[:, 4] - res[:, 2]) + (cdd[3] - cdd[1]) * (
|
|
cdd[4] - cdd[2]) - intersec_map)
|
|
res = res[iou_map_cur < iou_thresh]
|
|
|
|
return np.array(cdd_results)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
a = hard_nms(np.array([
|
|
[0.4, 1, 10, 12, 20],
|
|
[0.5, 1, 11, 11, 20],
|
|
[0.55, 20, 30, 40, 50]
|
|
]), topn=100, iou_thresh=0.4)
|
|
print(a)
|