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 []