def decode(self, loc_preds, cls_preds, input_size):
'''Decode outputs back to bouding box locations and class labels.
Args:
loc_preds: (tensor) predicted locations, sized [#anchors, 4].
cls_preds: (tensor) predicted class labels, sized [#anchors, #classes].
input_size: (int/tuple) model input size of (w,h).
Returns:
boxes: (tensor) decode box locations, sized [#obj,4].
labels: (tensor) class labels for each box, sized [#obj,].
'''
CLS_THRESH = 0.5
NMS_THRESH = 0.5
input_size = torch.Tensor([input_size,input_size]) if isinstance(input_size, int) \
else torch.Tensor(input_size)
anchor_boxes = self._get_anchor_boxes(input_size)
loc_xy = loc_preds[:,:2]
loc_wh = loc_preds[:,2:]
xy = loc_xy * anchor_boxes[:,2:] + anchor_boxes[:,:2]
wh = loc_wh.exp() * anchor_boxes[:,2:]
boxes = torch.cat([xy-wh/2, xy+wh/2], 1) # [#anchors,4]
score, labels = cls_preds.sigmoid().max(1) # [#anchors,]
ids = score > CLS_THRESH
ids = ids.nonzero().squeeze() # [#obj,]
if not ids.size() :
return torch.tensor([0, 0, 0, 0]), torch.tensor([0]), torch.tensor([0]), False
if ids.size()[0] == 0 :
return torch.tensor([0, 0, 0, 0]), torch.tensor([0]), torch.tensor([0]), False
keep, sco = box_nms(boxes[ids], score[ids], threshold=NMS_THRESH)
return boxes[ids][keep], labels[ids][keep], sco, True
def heatmap_to_nms(self, heatmap, tensor=False, boxnms=False):
"""
return:
heatmap_nms_batch: np [batch, 1, H, W]
"""
to_floatTensor = lambda x: torch.from_numpy(x).type(torch.FloatTensor)
from utils.var_dim import toNumpy
heatmap_np = toNumpy(heatmap)
## heatmap_nms
if boxnms:
from utils.utils import box_nms
heatmap_nms_batch = [box_nms(h.detach().squeeze(), self.nms_dist, min_prob=self.conf_thresh) \
for h in heatmap] # [batch, H, W]
heatmap_nms_batch = torch.stack(heatmap_nms_batch,
dim=0).unsqueeze(1)
# print('heatmap_nms_batch: ', heatmap_nms_batch.shape)
else:
heatmap_nms_batch = [self.heatmap_nms(h, self.nms_dist, self.conf_thresh) \
for h in heatmap_np] # [batch, H, W]
heatmap_nms_batch = np.stack(heatmap_nms_batch, axis=0)
heatmap_nms_batch = heatmap_nms_batch[:, np.newaxis, ...]
if tensor:
heatmap_nms_batch = to_floatTensor(heatmap_nms_batch)
heatmap_nms_batch = heatmap_nms_batch.to(self.device)
self.heatmap = heatmap
self.heatmap_nms_batch = heatmap_nms_batch
return heatmap_nms_batch
pass
def add2tensorboard_nms(self,
img,
labels_2D,
semi,
task="training",
batch_size=1):
"""
# deprecated:
:param img:
:param labels_2D:
:param semi:
:param task:
:param batch_size:
:return:
"""
from utils.utils import getPtsFromHeatmap
from utils.utils import box_nms
boxNms = False
n_iter = self.n_iter
nms_dist = self.config["model"]["nms"]
conf_thresh = self.config["model"]["detection_threshold"]
# print("nms_dist: ", nms_dist)
precision_recall_list = []
precision_recall_boxnms_list = []
for idx in range(batch_size):
semi_flat_tensor = flattenDetection(semi[idx, :, :, :]).detach()
semi_flat = toNumpy(semi_flat_tensor)
semi_thd = np.squeeze(semi_flat, 0)
pts_nms = getPtsFromHeatmap(semi_thd, conf_thresh, nms_dist)
semi_thd_nms_sample = np.zeros_like(semi_thd)
semi_thd_nms_sample[pts_nms[1, :].astype(np.int),
pts_nms[0, :].astype(np.int)] = 1
label_sample = torch.squeeze(labels_2D[idx, :, :, :])
# pts_nms = getPtsFromHeatmap(label_sample.numpy(), conf_thresh, nms_dist)
# label_sample_rms_sample = np.zeros_like(label_sample.numpy())
# label_sample_rms_sample[pts_nms[1, :].astype(np.int), pts_nms[0, :].astype(np.int)] = 1
label_sample_nms_sample = label_sample
if idx < 5:
result_overlap = img_overlap(
np.expand_dims(label_sample_nms_sample, 0),
np.expand_dims(semi_thd_nms_sample, 0),
toNumpy(img[idx, :, :, :]),
)
self.writer.add_image(
task + "-detector_output_thd_overlay-NMS" + "/%d" % idx,
result_overlap,
n_iter,
)
assert semi_thd_nms_sample.shape == label_sample_nms_sample.size()
precision_recall = precisionRecall_torch(
torch.from_numpy(semi_thd_nms_sample), label_sample_nms_sample)
precision_recall_list.append(precision_recall)
if boxNms:
semi_flat_tensor_nms = box_nms(semi_flat_tensor.squeeze(),
nms_dist,
min_prob=conf_thresh).cpu()
semi_flat_tensor_nms = (semi_flat_tensor_nms >=
conf_thresh).float()
if idx < 5:
result_overlap = img_overlap(
np.expand_dims(label_sample_nms_sample, 0),
semi_flat_tensor_nms.numpy()[np.newaxis, :, :],
toNumpy(img[idx, :, :, :]),
)
self.writer.add_image(
task + "-detector_output_thd_overlay-boxNMS" +
"/%d" % idx,
result_overlap,
n_iter,
)
precision_recall_boxnms = precisionRecall_torch(
semi_flat_tensor_nms, label_sample_nms_sample)
precision_recall_boxnms_list.append(precision_recall_boxnms)
precision = np.mean([
precision_recall["precision"]
for precision_recall in precision_recall_list
])
recall = np.mean([
precision_recall["recall"]
for precision_recall in precision_recall_list
])
self.writer.add_scalar(task + "-precision_nms", precision, n_iter)
self.writer.add_scalar(task + "-recall_nms", recall, n_iter)
print("-- [%s-%d-fast NMS] precision: %.4f, recall: %.4f" %
(task, n_iter, precision, recall))
if boxNms:
precision = np.mean([
precision_recall["precision"]
for precision_recall in precision_recall_boxnms_list
])
recall = np.mean([
precision_recall["recall"]
for precision_recall in precision_recall_boxnms_list
])
self.writer.add_scalar(task + "-precision_boxnms", precision,
n_iter)
self.writer.add_scalar(task + "-recall_boxnms", recall, n_iter)
print("-- [%s-%d-boxNMS] precision: %.4f, recall: %.4f" %
(task, n_iter, precision, recall))