def _compute_boxes(_score_map,
_geo_map): # 从得到的两个map中恢复出原图的检测框，测试时才使用，训练不用
score = _score_map.permute(0, 2, 3, 1)
geometry = _geo_map.permute(0, 2, 3, 1)
score = score.detach().cpu().numpy(
) # detach后放到cpu会导致张量计算图断开，无法反向传播
geometry = geometry.detach().cpu().numpy()
timer = {'net': 0, 'restore': 0, 'nms': 0}
_pred_mapping = [] # 下标为i的样本有几个检测框，用于做_pred_boxes的索引
_pred_boxes = [] # 所有预测出来的检测框
for i in range(score.shape[0]):
cur_score = score[i, :, :, 0]
cur_geometry = geometry[i, :, :, ]
detected_boxes, _ = Toolbox.detect(score_map=cur_score,
geo_map=cur_geometry,
timer=timer)
if detected_boxes is None:
continue
num_detected_boxes = detected_boxes.shape[0]
if len(detected_boxes) > 0:
_pred_mapping.append(np.array([i] * num_detected_boxes))
_pred_boxes.append(detected_boxes)
return np.concatenate(_pred_boxes) if len(_pred_boxes) > 0 else [], \
np.concatenate(_pred_mapping) if len(_pred_mapping) > 0 else []
def call(self, input):
image, boxes, mapping = input
feature_map = self.sharedConv(image)
score_map, geo_map = self.detector(feature_map)
if self.training:
rois, lengths, indices = self.roirotate.call(
feature_map, boxes[:, :8], mapping)
pred_mapping = mapping
pred_boxes = boxes
else:
score = score_map.permute(0, 2, 3, 1)
geometry = geo_map.permute(0, 2, 3, 1)
score = score.detach().cpu().numpy()
geometry = geometry.detach().cpu().numpy()
timer = {'net': 0, 'restore': 0, 'nms': 0}
pred_boxes = []
pred_mapping = []
for i in range(score.shape[0]):
s = score[i, :, :, 0]
g = geometry[i, :, :, ]
bb, _ = Toolbox.detect(score_map=s, geo_map=g, timer=timer)
bb_size = bb.shape[0]
if len(bb) > 0:
pred_mapping.append(np.array([i] * bb_size))
pred_boxes.append(bb)
if len(pred_mapping) > 0:
pred_boxes = np.concatenate(pred_boxes)
pred_mapping = np.concatenate(pred_mapping)
rois, lengths, indices = self.roirotate.call(
feature_map, pred_boxes[:, :8], pred_mapping)
else:
return score_map, geo_map, (
None, None), pred_boxes, pred_mapping, None
lengths = tf.convert_to_tensor(lengths)
preds = self.recognizer(rois, lengths)
preds = preds.permute(1, 0, 2) # B, T, C -> T, B, C
return score_map, geo_map, (preds,
lengths), pred_boxes, pred_mapping, indices
def _compute_boxes(_score_map, _geo_map):
score = _score_map.permute(0, 2, 3, 1)
geometry = _geo_map.permute(0, 2, 3, 1)
score = score.detach().cpu().numpy()
geometry = geometry.detach().cpu().numpy()
timer = {'net': 0, 'restore': 0, 'nms': 0}
_pred_mapping = []
_pred_boxes = []
for i in range(score.shape[0]):
cur_score = score[i, :, :, 0]
cur_geometry = geometry[i, :, :, ]
detected_boxes, _ = Toolbox.detect(score_map=cur_score,
geo_map=cur_geometry,
timer=timer)
if detected_boxes is None:
continue
num_detected_boxes = detected_boxes.shape[0]
if len(detected_boxes) > 0:
_pred_mapping.append(np.array([i] * num_detected_boxes))
_pred_boxes.append(detected_boxes)
return np.concatenate(_pred_boxes) if len(_pred_boxes) > 0 else [], \
np.concatenate(_pred_mapping) if len(_pred_mapping) > 0 else []