|
|
import torch |
|
|
|
|
|
from mmdet.core import multi_apply |
|
|
from ..builder import HEADS |
|
|
from ..losses import CrossEntropyLoss, SmoothL1Loss, carl_loss, isr_p |
|
|
from .ssd_head import SSDHead |
|
|
|
|
|
|
|
|
|
|
|
@HEADS.register_module() |
|
|
class PISASSDHead(SSDHead): |
|
|
|
|
|
def loss(self, |
|
|
cls_scores, |
|
|
bbox_preds, |
|
|
gt_bboxes, |
|
|
gt_labels, |
|
|
img_metas, |
|
|
gt_bboxes_ignore=None): |
|
|
"""Compute losses of the head. |
|
|
|
|
|
Args: |
|
|
cls_scores (list[Tensor]): Box scores for each scale level |
|
|
Has shape (N, num_anchors * num_classes, H, W) |
|
|
bbox_preds (list[Tensor]): Box energies / deltas for each scale |
|
|
level with shape (N, num_anchors * 4, H, W) |
|
|
gt_bboxes (list[Tensor]): Ground truth bboxes of each image |
|
|
with shape (num_obj, 4). |
|
|
gt_labels (list[Tensor]): Ground truth labels of each image |
|
|
with shape (num_obj, 4). |
|
|
img_metas (list[dict]): Meta information of each image, e.g., |
|
|
image size, scaling factor, etc. |
|
|
gt_bboxes_ignore (list[Tensor]): Ignored gt bboxes of each image. |
|
|
Default: None. |
|
|
|
|
|
Returns: |
|
|
dict: Loss dict, comprise classification loss regression loss and |
|
|
carl loss. |
|
|
""" |
|
|
featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] |
|
|
assert len(featmap_sizes) == self.anchor_generator.num_levels |
|
|
|
|
|
device = cls_scores[0].device |
|
|
|
|
|
anchor_list, valid_flag_list = self.get_anchors( |
|
|
featmap_sizes, img_metas, device=device) |
|
|
cls_reg_targets = self.get_targets( |
|
|
anchor_list, |
|
|
valid_flag_list, |
|
|
gt_bboxes, |
|
|
img_metas, |
|
|
gt_bboxes_ignore_list=gt_bboxes_ignore, |
|
|
gt_labels_list=gt_labels, |
|
|
label_channels=1, |
|
|
unmap_outputs=False, |
|
|
return_sampling_results=True) |
|
|
if cls_reg_targets is None: |
|
|
return None |
|
|
(labels_list, label_weights_list, bbox_targets_list, bbox_weights_list, |
|
|
num_total_pos, num_total_neg, sampling_results_list) = cls_reg_targets |
|
|
|
|
|
num_images = len(img_metas) |
|
|
all_cls_scores = torch.cat([ |
|
|
s.permute(0, 2, 3, 1).reshape( |
|
|
num_images, -1, self.cls_out_channels) for s in cls_scores |
|
|
], 1) |
|
|
all_labels = torch.cat(labels_list, -1).view(num_images, -1) |
|
|
all_label_weights = torch.cat(label_weights_list, |
|
|
-1).view(num_images, -1) |
|
|
all_bbox_preds = torch.cat([ |
|
|
b.permute(0, 2, 3, 1).reshape(num_images, -1, 4) |
|
|
for b in bbox_preds |
|
|
], -2) |
|
|
all_bbox_targets = torch.cat(bbox_targets_list, |
|
|
-2).view(num_images, -1, 4) |
|
|
all_bbox_weights = torch.cat(bbox_weights_list, |
|
|
-2).view(num_images, -1, 4) |
|
|
|
|
|
|
|
|
all_anchors = [] |
|
|
for i in range(num_images): |
|
|
all_anchors.append(torch.cat(anchor_list[i])) |
|
|
|
|
|
isr_cfg = self.train_cfg.get('isr', None) |
|
|
all_targets = (all_labels.view(-1), all_label_weights.view(-1), |
|
|
all_bbox_targets.view(-1, |
|
|
4), all_bbox_weights.view(-1, 4)) |
|
|
|
|
|
if isr_cfg is not None: |
|
|
all_targets = isr_p( |
|
|
all_cls_scores.view(-1, all_cls_scores.size(-1)), |
|
|
all_bbox_preds.view(-1, 4), |
|
|
all_targets, |
|
|
torch.cat(all_anchors), |
|
|
sampling_results_list, |
|
|
loss_cls=CrossEntropyLoss(), |
|
|
bbox_coder=self.bbox_coder, |
|
|
**self.train_cfg.isr, |
|
|
num_class=self.num_classes) |
|
|
(new_labels, new_label_weights, new_bbox_targets, |
|
|
new_bbox_weights) = all_targets |
|
|
all_labels = new_labels.view(all_labels.shape) |
|
|
all_label_weights = new_label_weights.view(all_label_weights.shape) |
|
|
all_bbox_targets = new_bbox_targets.view(all_bbox_targets.shape) |
|
|
all_bbox_weights = new_bbox_weights.view(all_bbox_weights.shape) |
|
|
|
|
|
|
|
|
carl_loss_cfg = self.train_cfg.get('carl', None) |
|
|
if carl_loss_cfg is not None: |
|
|
loss_carl = carl_loss( |
|
|
all_cls_scores.view(-1, all_cls_scores.size(-1)), |
|
|
all_targets[0], |
|
|
all_bbox_preds.view(-1, 4), |
|
|
all_targets[2], |
|
|
SmoothL1Loss(beta=1.), |
|
|
**self.train_cfg.carl, |
|
|
avg_factor=num_total_pos, |
|
|
num_class=self.num_classes) |
|
|
|
|
|
|
|
|
assert torch.isfinite(all_cls_scores).all().item(), \ |
|
|
'classification scores become infinite or NaN!' |
|
|
assert torch.isfinite(all_bbox_preds).all().item(), \ |
|
|
'bbox predications become infinite or NaN!' |
|
|
|
|
|
losses_cls, losses_bbox = multi_apply( |
|
|
self.loss_single, |
|
|
all_cls_scores, |
|
|
all_bbox_preds, |
|
|
all_anchors, |
|
|
all_labels, |
|
|
all_label_weights, |
|
|
all_bbox_targets, |
|
|
all_bbox_weights, |
|
|
num_total_samples=num_total_pos) |
|
|
loss_dict = dict(loss_cls=losses_cls, loss_bbox=losses_bbox) |
|
|
if carl_loss_cfg is not None: |
|
|
loss_dict.update(loss_carl) |
|
|
return loss_dict |
|
|
|
