未验证 提交 2ecc6525 编写于 作者: G George Ni 提交者: GitHub

[MOT] fix jde tensorrt deploy (#3617)

上级 c0cb4407
......@@ -40,7 +40,7 @@ TRT_MIN_SUBGRAPH = {
'HigherHRNet': 3,
'HRNet': 3,
'DeepSORT': 3,
'JDE': 3,
'JDE': 10,
'FairMOT': 5,
}
......
......@@ -927,37 +927,41 @@ class JDEBox(object):
gy2 = gy + gh * 0.5
return paddle.stack([gx1, gy1, gx2, gy2], axis=1)
def decode_delta_map(self, delta_map, anchors):
nB, nA, nGh, nGw, _ = delta_map.shape[:]
anchor_mesh = self.generate_anchor(nGh, nGw, anchors)
# only support bs=1
def decode_delta_map(self, nA, nGh, nGw, delta_map, anchor_vec):
anchor_mesh = self.generate_anchor(nGh, nGw, anchor_vec)
anchor_mesh = paddle.unsqueeze(anchor_mesh, 0)
pred_list = self.decode_delta(
paddle.reshape(
delta_map, shape=[-1, 4]),
paddle.reshape(
anchor_mesh, shape=[-1, 4]))
pred_map = paddle.reshape(pred_list, shape=[nB, nA * nGh * nGw, 4])
pred_map = paddle.reshape(pred_list, shape=[nA * nGh * nGw, 4])
return pred_map
def _postprocessing_by_level(self, nA, stride, head_out, anchor_vec):
boxes_shape = head_out.shape
nB, nGh, nGw = 1, boxes_shape[-2], boxes_shape[-1]
# only support bs=1
boxes_shape = head_out.shape # [nB, nA*6, nGh, nGw]
nGh, nGw = boxes_shape[-2], boxes_shape[-1]
nB = 1 # TODO: only support bs=1 now
boxes_list, scores_list = [], []
for idx in range(nB):
p = paddle.reshape(
head_out, shape=[nB, nA, self.num_classes + 5, nGh, nGw])
p = paddle.transpose(p, perm=[0, 1, 3, 4, 2]) # [nB, 4, nGh, nGw, 6]
p_box = p[:, :, :, :, :4]
boxes = self.decode_delta_map(p_box, anchor_vec) # [nB, 4*nGh*nGw, 4]
boxes = boxes * stride
head_out[idx], shape=[nA, self.num_classes + 5, nGh, nGw])
p = paddle.transpose(p, perm=[0, 2, 3, 1]) # [nA, nGh, nGw, 6]
delta_map = p[:, :, :, :4]
boxes = self.decode_delta_map(nA, nGh, nGw, delta_map, anchor_vec)
# [nA * nGh * nGw, 4]
boxes_list.append(boxes * stride)
p_conf = paddle.transpose(
p[:, :, :, :, 4:6], perm=[0, 4, 1, 2, 3]) # [nB, 2, 4, 19, 34]
p[:, :, :, 4:6], perm=[3, 0, 1, 2]) # [2, nA, nGh, nGw]
p_conf = F.softmax(
p_conf, axis=1)[:, 1, :, :, :].unsqueeze(-1) # [nB, 4, 19, 34, 1]
scores = paddle.reshape(p_conf, shape=[nB, nA * nGh * nGw, 1])
return boxes, scores
p_conf, axis=0)[1, :, :, :].unsqueeze(-1) # [nA, nGh, nGw, 1]
scores = paddle.reshape(p_conf, shape=[nA * nGh * nGw, 1])
scores_list.append(scores)
boxes_results = paddle.stack(boxes_list)
scores_results = paddle.stack(scores_list)
return boxes_results, scores_results
def __call__(self, yolo_head_out, anchors):
bbox_pred_list = []
......
......@@ -175,8 +175,7 @@ class JDEEmbeddingHead(nn.Layer):
for i, p_ide in enumerate(ide_outs):
p_ide = p_ide.transpose((0, 2, 3, 1))
p_ide_repeat = paddle.tile(
p_ide.unsqueeze(axis=0), [1, self.anchor_scales, 1, 1, 1])
p_ide_repeat = paddle.tile(p_ide, [self.anchor_scales, 1, 1, 1])
embedding = F.normalize(p_ide_repeat, axis=-1)
emb = paddle.reshape(embedding, [-1, self.embedding_dim])
emb_outs.append(emb)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册