def nms(dets, thresh, force_cpu=False):
"""Dispatch to either CPU or GPU NMS implementations."""
if dets.shape[0] == 0:
return []
# ---numpy version---
# original: return gpu_nms(dets, thresh, device_id=cfg.GPU_ID)
# ---pytorch version---
# Force nms_gpu Compile check
try:
return nms_gpu(dets, thresh) if force_cpu == False else nms_cpu(
dets, thresh)
except NameError:
return nms_cpu(dets, thresh)
def nms(dets, thresh, force_cpu=False):
"""Dispatch to either CPU or GPU NMS implementations."""
if dets.shape[0] == 0:
return []
# ---numpy version---
# original: return gpu_nms(dets, thresh, device_id=cfg.GPU_ID)
# ---pytorch version---
return nms_gpu(dets, thresh) if force_cpu == False else nms_cpu(dets, thresh)
def forward(self, loc_data, conf_data, prior_data):
bs = loc_data.size(0) # batch size
num_priors = prior_data.size(0)
output = torch.zeros(bs, self.num_classes * self.top_k,
6).to(loc_data.device)
conf_preds = \
conf_data.view(bs, num_priors, self.num_classes).transpose(2, 1)
# Decode predictions into bboxes.
for i in range(bs):
decoded_boxes = decode(loc_data[i], prior_data, self.variance)
# For each class, perform nms
conf_scores = conf_preds[i].clone()
det_max = []
for cl in range(1, self.num_classes):
c_mask = conf_scores[cl].gt(self.conf_thresh)
scores = conf_scores[cl][c_mask].view(-1)
if scores.size(0) == 0:
continue
l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)
boxes = decoded_boxes[l_mask].view(-1, 4)
# idx of highest scoring and non-overlapping boxes per class
keep = nms_cpu(boxes, scores, self.nms_thresh, self.top_k)
count = len(keep)
if count > 0:
_classes = torch.tensor([[cl * 1.0] for _ in range(count)
]).to(loc_data.device)
det_max.append(
torch.cat(
(boxes[keep], scores[keep].unsqueeze(1), _classes),
1))
if len(det_max):
det_max = torch.cat(det_max)
output[i, :len(det_max)] = det_max[det_max[:, 4].argsort(
descending=True)]
return output
if vis:
im = cv2.imread(imdb.image_path_at(i))
im2show = np.copy(im)
for j in xrange(1, imdb.num_obj_classes):
inds = torch.nonzero(scores[:, j] >= thresh).view(-1)
# if there is det
if inds.numel() > 0:
cls_scores = scores[:, j][inds.data]
_, order = torch.sort(cls_scores, 0, True)
cls_boxes = torch.FloatTensor(
np.array(pred_boxes)[inds.data, j, :])
cls_dets = torch.cat(
(cls_boxes, cls_scores.data.cpu().unsqueeze(1)), 1)
cls_dets = cls_dets[order.data.cpu()]
keep = nms_cpu(cls_dets, cfg.TEST.NMS)
cls_dets = cls_dets[keep.view(-1).long()]
if vis:
im2show = vis_detections(im2show, imdb._obj_classes[j],
cls_dets.cpu().numpy(), 0.3)
all_boxes[j][i] = cls_dets.cpu().numpy()
else:
all_boxes[j][i] = empty_array
# Limit to max_per_image detections *over all classes*
if max_per_image > 0:
image_scores = np.hstack([
all_boxes[j][i][:, -1] for j in xrange(1, imdb.num_obj_classes)
])
if len(image_scores) > max_per_image:
image_thresh = np.sort(image_scores)[-max_per_image]