def nms(dets, thresh):
return nms_torch(dets[:, :4], dets[:, 4], thresh)
def nms(det, threshold):
return nms_torch(det[:, :4], det[:, 4], threshold)
def nms(dets, thresh):
#print("a", dets[:, :4].size())
#print("b", dets[:, 4].size())
return nms_torch(dets[:, :4], dets[:, 4], thresh)
def forward(self, inputs, targets=None, meta=None):
bone_feats = self.backbone(inputs)
trans_feats = []
for feat, conv in zip(bone_feats, self.trans_conv):
trans_feats.append(conv(feat))
features = self.bifpn(trans_feats)
regression = torch.cat(
[self.regressor(feature) for feature in features], dim=1)
classification = torch.cat(
[self.classifier(feature) for feature in features], dim=1)
anchors = self.anchors(inputs)
if self.training:
cls_loss, reg_loss = self.focalLoss(classification, regression,
anchors, targets)
return {
'classification_loss': cls_loss,
'regression_loss': reg_loss
}
else:
# 增加batch预测功能
transformed_anchors = self.regressBoxes(anchors, regression)
# anchor预测结果 (B, N, 4)
transformed_anchors = self.clipBoxes(transformed_anchors, inputs)
# 每个anchor对应的类别的最大得分
scores = torch.max(classification, dim=2, keepdim=True)[0]
# 通过得分阈值筛选有效anchor位置
scores_over_thresh = (scores > self.score_thr)[:, :, 0]
# 对每一张图像分别进行后处理
batch_size = scores.shape[0]
if meta is not None:
scales = meta['scale']
outputs = []
for i in range(batch_size):
result = []
# 第i个样本的有效anchor位置
scores_over_thresh_i = scores_over_thresh[i, :]
# 有矩形框进行后处理
if scores_over_thresh_i.sum() > 0:
# 第i个样本的每个类别得分
classification_i = classification[i,
scores_over_thresh_i, :]
# 第i个样本的举行框
transformed_anchors_i = transformed_anchors[
i, scores_over_thresh_i, :]
# 第i个样本的所有有效anchor分数
scores_i = scores[i, scores_over_thresh_i, 0]
# 通过分数进行nms（所有类别，此处不是multi-label nms）
anchors_nms_idx = nms_torch(transformed_anchors_i,
scores_i, self.iou_thr)
# 获取通过nms后保留下来的anchor分数以及类别
nms_scores_i, nms_class_i = classification_i[
anchors_nms_idx, :].max(dim=1)
# nms后保留下来的矩形框
nms_bboxes_i = transformed_anchors_i[anchors_nms_idx, :]
# xyxy -> xywh
nms_bboxes_i[:, 2] -= nms_bboxes_i[:, 0]
nms_bboxes_i[:, 3] -= nms_bboxes_i[:, 1]
# scale to the original size
nms_bboxes_i /= scales[i]
for j in range(nms_bboxes_i.shape[0]):
_score = float(nms_scores_i[j])
_label = int(nms_class_i[j])
_box = nms_bboxes_i[j, :]
item = {
'class': _label,
'score': _score,
'bbox': _box.tolist(),
}
result.append(item)
outputs.append(result)
return outputs