def forward(self, is_train, req, in_data, out_data, aux):
if self.nms_type == 'nms':
nms = py_nms_wrapper(self.nms_thr)
else:
raise NotImplementedError
cls_score = in_data[0].asnumpy()
bbox_xyxy = in_data[1].asnumpy()
cls_score_shape = cls_score.shape # (b, n, num_class_withbg)
bbox_xyxy_shape = bbox_xyxy.shape # (b, n, 4) or (b, n, 4 * num_class_withbg)
batch_image = cls_score_shape[0]
num_bbox = cls_score_shape[1]
num_class_withbg = cls_score_shape[2]
post_score = np.zeros((batch_image, self.max_det_per_image, 1), dtype=np.float32)
post_bbox_xyxy = np.zeros((batch_image, self.max_det_per_image, 4), dtype=np.float32)
post_cls = np.full((batch_image, self.max_det_per_image, 1), -1, dtype=np.float32)
for i, (per_image_cls_score, per_image_bbox_xyxy) in enumerate(zip(cls_score, bbox_xyxy)):
cls_det = multiclass_nms(nms, per_image_cls_score, per_image_bbox_xyxy, \
self.min_det_score, self.max_det_per_image)
num_det = cls_det.shape[0]
post_bbox_xyxy[i, :num_det] = cls_det[:, :4]
post_score[i, :num_det] = cls_det[:, -2][:, np.newaxis] # convert to (n, 1)
post_cls[i, :num_det] = cls_det[:, -1][:, np.newaxis] # convert to (n, 1)
self.assign(out_data[0], req[0], post_score)
self.assign(out_data[1], req[1], post_bbox_xyxy)
self.assign(out_data[2], req[2], post_cls)
def __init__(self, configFn, ctx, outFolder, threshold):
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
config = importlib.import_module(
configFn.replace('.py', '').replace('/', '.'))
_, _, _, _, _, _, self.__pModel, _, self.__pTest, self.transform, _, _, _ = config.get_config(
is_train=False)
self.__pModel = patch_config_as_nothrow(self.__pModel)
self.__pTest = patch_config_as_nothrow(self.__pTest)
self.resizeParam = (800, 1200)
if callable(self.__pTest.nms.type):
self.__nms = self.__pTest.nms.type(self.__pTest.nms.thr)
else:
from operator_py.nms import py_nms_wrapper
self.__nms = py_nms_wrapper(self.__pTest.nms.thr)
arg_params, aux_params = load_checkpoint(self.__pTest.model.prefix,
self.__pTest.model.epoch)
sym = self.__pModel.test_symbol
from utils.graph_optimize import merge_bn
sym, arg_params, aux_params = merge_bn(sym, arg_params, aux_params)
self.__mod = DetModule(
sym,
data_names=['data', 'im_info', 'im_id', 'rec_id'],
context=ctx)
self.__mod.bind(data_shapes=[('data', (1, 3, self.resizeParam[0],
self.resizeParam[1])),
('im_info', (1, 3)), ('im_id', (1, )),
('rec_id', (1, ))],
for_training=False)
self.__mod.set_params(arg_params, aux_params, allow_extra=False)
self.__saveSymbol(sym, outFolder,
self.__pTest.model.prefix.split('/')[-1])
self.__threshold = threshold
self.outFolder = outFolder
def __init__(self, config, batch_size, gpu_id, thresh):
self.config = config
self.batch_size = batch_size
self.thresh = thresh
# Parse the parameter file of model
pGen, pKv, pRpn, pRoi, pBbox, pDataset, pModel, pOpt, pTest, \
transform, data_name, label_name, metric_list = config.get_config(is_train=False)
self.data_name = data_name
self.label_name = label_name
self.p_long, self.p_short = transform[1].p.long, transform[1].p.short
# Define NMS type
if callable(pTest.nms.type):
self.do_nms = pTest.nms.type(pTest.nms.thr)
else:
from operator_py.nms import py_nms_wrapper
self.do_nms = py_nms_wrapper(pTest.nms.thr)
sym = pModel.test_symbol
sym.save(pTest.model.prefix + "_test.json")
ctx = mx.gpu(gpu_id)
data_shape = [
('data', (batch_size, 3, 800, 1200)),
("im_info", (1, 3)),
("im_id", (1, )),
("rec_id", (1, )),
]
# Load network
arg_params, aux_params = load_checkpoint(pTest.model.prefix,
pTest.model.epoch)
self.mod = DetModule(sym, data_names=data_name, context=ctx)
self.mod.bind(data_shapes=data_shape, for_training=False)
self.mod.set_params(arg_params, aux_params, allow_extra=False)
output_dict[k]["bbox_xyxy"] = output_dict[k]["bbox_xyxy"][0]
if len(output_dict[k]["cls_score"]) > 1:
output_dict[k]["cls_score"] = np.concatenate(
output_dict[k]["cls_score"])
else:
output_dict[k]["cls_score"] = output_dict[k]["cls_score"][0]
t3_s = time.time()
print("aggregate uses: %.1f" % (t3_s - t2_s))
if callable(pTest.nms.type):
nms = pTest.nms.type(pTest.nms.thr)
else:
from operator_py.nms import py_nms_wrapper
nms = py_nms_wrapper(pTest.nms.thr)
def do_nms(k):
bbox_xyxy = output_dict[k]["bbox_xyxy"]
cls_score = output_dict[k]["cls_score"]
final_dets = {}
for cid in range(cls_score.shape[1]):
score = cls_score[:, cid]
if bbox_xyxy.shape[1] != 4:
cls_box = bbox_xyxy[:, cid * 4:(cid + 1) * 4]
else:
cls_box = bbox_xyxy
valid_inds = np.where(score > pTest.min_det_score)[0]
box = cls_box[valid_inds]
score = score[valid_inds]