def compute_loss(self, x, targets):
xs = self.feature_extractor(x)
ys = self.feature_extractor2(xs)
loc_preds, cls_preds = self.rpn(xs)
loc_preds2, cls_preds2 = self.rpn(xs)
with torch.no_grad():
loc_targets, cls_targets = self.box_coder.encode(xs, targets)
with torch.no_grad():
anchors, anchors_xyxy = self.box_coder(xs)
anchors2 = box.deltas_to_bbox(loc_preds, anchors)
anchors2xyxy = box.change_box_order(anchors2, 'xywh2xyxy')
loc_targets2, cls_targets2 = self.box_coder.encode_with_anchors(
anchors2, anchors2xyxy, targets)
#cls_targets *= cls_preds?
assert cls_targets.shape[1] == cls_preds.shape[1]
loc_loss, cls_loss = self.criterion(loc_preds, loc_targets, cls_preds,
cls_targets)
loc_loss2, cls_loss2 = self.criterion(loc_preds2, loc_targets2,
cls_preds2, cls_targets2)
loss_dict = {
'loc': loc_loss,
'cls_loss': cls_loss,
'loc2': loc_loss2,
'cls_loss2': cls_loss2
}
return loss_dict
def reset(self, ssd_model):
self.steps = ssd_model.steps
self.box_sizes = ssd_model.box_sizes
self.aspect_ratios = ssd_model.aspect_ratios
self.fm_sizes = ssd_model.fm_sizes
self.height = ssd_model.height
self.width = ssd_model.width
self.default_boxes = self._get_default_boxes()
self.default_boxes_xyxy = change_box_order(self.default_boxes,
'xywh2xyxy')
def reset(self, ssd_model):
self.steps = ssd_model.steps
self.box_sizes = ssd_model.box_sizes
self.aspect_ratios = ssd_model.aspect_ratios
self.fm_sizes = ssd_model.fm_sizes
self.height = ssd_model.height
self.width = ssd_model.width
self.fm_len = []
self.register_buffer('default_boxes', self._get_default_boxes())
self.register_buffer('default_boxes_xyxy',
change_box_order(self.default_boxes, 'xywh2xyxy'))
def encode_fast(self, gt_boxes, labels):
boxes, cls_targets = assign_priors(gt_boxes, labels + 1,
self.default_boxes_xyxy,
self.fg_iou_threshold,
self.bg_iou_threshold)
boxes = change_box_order(boxes, 'xyxy2xywh')
default_boxes = self.default_boxes
loc_xy = (boxes[:, :2] - default_boxes[:, :2]
) / default_boxes[:, 2:] / self.variances[0]
loc_wh = torch.log(
boxes[:, 2:] / default_boxes[:, 2:]) / self.variances[1]
loc_targets = torch.cat([loc_xy, loc_wh], 1)
return loc_targets, cls_targets
def pytestcase_all_gt_should_be_matched_even_low_iou(self):
"""
Boxes with small iou should be matched
:return:
"""
box_coder = Anchors(allow_low_quality_matches=True)
anchors_xyxy = torch.tensor(
[[25, 25, 250, 250], [20, 20, 50, 50], [3, 3, 4, 4]],
dtype=torch.float32)
anchors = box.change_box_order(anchors_xyxy, 'xyxy2xywh')
targets = torch.tensor([[120, 120, 250, 250, 1], [20, 20, 22, 22, 2]],
dtype=torch.float32)
targets = [[targets]]
_, cls_targets = box_coder.encode(anchors, anchors_xyxy, targets)
assert len(torch.unique(
cls_targets)) == 3 # first box and +1 is for background class
def gather_boxes(self, proposals):
#this expects list of tensor of shape N, 4
idxs = []
sizes = []
rois = []
rois_xyxy = []
stride = len(proposals)
for t in range(len(proposals)):
for i in range(len(proposals[t])):
boxes, _, _ = proposals[t][i]
num = len(boxes) if boxes is not None else 0
sizes += [num]
if num > 0:
boxes = boxes.detach()
rois_xyxy += [boxes]
rois += [box.change_box_order(boxes, 'xyxy2xywh')]
idxs += [t * stride + i] * num
idxs = torch.LongTensor(idxs).to(rois[0].device)
return rois, rois_xyxy, sizes, idxs
def __init__(self,
ssd_model,
fg_iou_threshold=0.6,
bg_iou_threshold=0.4,
soft_nms=False):
super(SSDBoxCoder, self).__init__()
self.steps = ssd_model.steps
self.box_sizes = ssd_model.box_sizes
self.aspect_ratios = ssd_model.aspect_ratios
self.scales = ssd_model.scales
self.fm_sizes = ssd_model.fm_sizes
self.height = ssd_model.height
self.width = ssd_model.width
self.fm_len = []
self.register_buffer('default_boxes', self._get_default_boxes_v2())
self.register_buffer('default_boxes_xyxy',
change_box_order(self.default_boxes, 'xywh2xyxy'))
self.fg_iou_threshold = fg_iou_threshold
self.bg_iou_threshold = bg_iou_threshold
self.use_cuda = False
self.variances = (0.1, 0.2)
self.nms = box_soft_nms if soft_nms else box_nms
self.encode = self.encode_fast
def encode_alt(self, boxes, gt_labels):
labels = gt_labels + 1
default_boxes = self.default_boxes_xyxy
ious = box_iou(default_boxes, boxes) # [#anchors, #obj]
index = torch.zeros(len(default_boxes),
dtype=torch.int64,
device=boxes.device).fill_(-1)
# We match every ground truth with higher than iou_threshold iou with an anchor
max_iou_anchors, arg_max_iou_anchors = torch.max(ious, dim=1)
mask_greater_than_threshold = max_iou_anchors >= self.fg_iou_threshold
mask_ambiguous = (max_iou_anchors >= self.bg_iou_threshold) * (
max_iou_anchors < self.fg_iou_threshold)
index[mask_greater_than_threshold] = arg_max_iou_anchors[
mask_greater_than_threshold]
index[mask_ambiguous] = -1
# We match every ground truth to an anchor the maximum iou iteratively (avoid two gt matched with same anchor)
for index_gt in range(len(boxes)):
arg_max_iou_gt = torch.argmax(ious[:, index_gt])
if ious[arg_max_iou_gt, index_gt] >= 0.2:
index[arg_max_iou_gt] = index_gt
ious[arg_max_iou_gt, :] = 0
boxes = boxes[index.clamp(min=0)] # negative index not supported
boxes = change_box_order(boxes, 'xyxy2xywh')
default_boxes = self.default_boxes # change_box_order(default_boxes, 'xyxy2xywh')
variances = (0.1, 0.2)
loc_xy = (boxes[:, :2] -
default_boxes[:, :2]) / default_boxes[:, 2:] / variances[0]
loc_wh = torch.log(boxes[:, 2:] / default_boxes[:, 2:]) / variances[1]
loc_targets = torch.cat([loc_xy, loc_wh], 1)
cls_targets = labels[index.clamp(min=0)]
cls_targets[index < 0] = 0
return loc_targets, cls_targets
def get_refined_anchors(self, xs, loc_preds):
anchors, _ = self.box_coder(xs)
anchors2 = box.deltas_to_bbox(loc_preds, anchors)
anchors2xyxy = box.change_box_order(anchors2, 'xywh2xyxy')
return anchors2, anchors2xyxy
