def forward_for_single_feature_map(self, box_cls, box_regression,
centerness, anchors):
N, _, H, W = box_cls.shape
A = box_regression.size(1) // 4
C = box_cls.size(1) // A
# put in the same format as anchors
box_cls = permute_and_flatten(box_cls, N, A, C, H, W)
box_cls = box_cls.sigmoid()
box_regression = permute_and_flatten(box_regression, N, A, 4, H, W)
box_regression = box_regression.reshape(N, -1, 4)
candidate_inds = box_cls > self.pre_nms_thresh
pre_nms_top_n = candidate_inds.view(N, -1).sum(1)
pre_nms_top_n = pre_nms_top_n.clamp(max=self.pre_nms_top_n)
centerness = permute_and_flatten(centerness, N, A, 1, H, W)
centerness = centerness.reshape(N, -1).sigmoid()
# multiply the classification scores with centerness scores
box_cls = box_cls * centerness[:, :, None]
results = []
for per_box_cls, per_box_regression, per_pre_nms_top_n, per_candidate_inds, per_anchors \
in zip(box_cls, box_regression, pre_nms_top_n, candidate_inds, anchors):
per_box_cls = per_box_cls[per_candidate_inds]
per_box_cls, top_k_indices = per_box_cls.topk(per_pre_nms_top_n,
sorted=False)
per_candidate_nonzeros = per_candidate_inds.nonzero()[
top_k_indices, :]
per_box_loc = per_candidate_nonzeros[:, 0]
per_class = per_candidate_nonzeros[:, 1] + 1
detections = self.box_coder.decode(
per_box_regression[per_box_loc, :].view(-1, 4),
per_anchors.bbox[per_box_loc, :].view(-1, 4))
boxlist = BoxList(detections, per_anchors.size, mode="xyxy")
boxlist.add_field("labels", per_class)
boxlist.add_field("scores", torch.sqrt(per_box_cls))
boxlist = boxlist.clip_to_image(remove_empty=False)
boxlist = remove_small_boxes(boxlist, self.min_size)
results.append(boxlist)
return results
def __getitem__(self, index):
img, annot = super().__getitem__(index)
annot = [o for o in annot if o['iscrowd'] == 0]
boxes = [o['bbox'] for o in annot]
boxes = torch.as_tensor(boxes).reshape(-1, 4)
target = BoxList(boxes, img.size, mode='xywh').convert('xyxy')
classes = [o['category_id'] for o in annot]
classes = [self.category2id[c] for c in classes]
classes = torch.tensor(classes)
# target.fields['labels'] = classes
target.add_field('labels', classes)
target.clip_to_image(remove_empty=True)
if self.transformer is not None:
img, target = self.transformer(img, target)
return img, target, index
