mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
37 lines
1.1 KiB
Python
Executable File
37 lines
1.1 KiB
Python
Executable File
from rcnn.config import config
|
|
import numpy as np
|
|
|
|
|
|
def compute_assign_targets(rois, threshold):
|
|
rois_area = np.sqrt(
|
|
(rois[:, 2] - rois[:, 0] + 1) * (rois[:, 3] - rois[:, 1] + 1))
|
|
num_rois = np.shape(rois)[0]
|
|
assign_levels = np.zeros(num_rois, dtype=np.uint8)
|
|
for i, stride in enumerate(config.RCNN_FEAT_STRIDE):
|
|
thd = threshold[i]
|
|
idx = np.logical_and(thd[1] <= rois_area, rois_area < thd[0])
|
|
assign_levels[idx] = stride
|
|
|
|
assert 0 not in assign_levels, "All rois should assign to specify levels."
|
|
return assign_levels
|
|
|
|
|
|
def add_assign_targets(roidb):
|
|
"""
|
|
given roidb, add ['assign_level']
|
|
:param roidb: roidb to be processed. must have gone through imdb.prepare_roidb
|
|
"""
|
|
print 'add assign targets'
|
|
assert len(roidb) > 0
|
|
assert 'boxes' in roidb[0]
|
|
|
|
area_threshold = [[np.inf, 448], [448, 224], [224, 112], [112, 0]]
|
|
|
|
assert len(config.RCNN_FEAT_STRIDE) == len(area_threshold)
|
|
|
|
num_images = len(roidb)
|
|
for im_i in range(num_images):
|
|
rois = roidb[im_i]['boxes']
|
|
roidb[im_i]['assign_levels'] = compute_assign_targets(
|
|
rois, area_threshold)
|