def heatmap_nms(heatmap, nms_dist=4, conf_thresh=0.015):
"""
input:
heatmap: np [(1), H, W]
"""
from utils.utils import getPtsFromHeatmap
# nms_dist = self.config['model']['nms']
# conf_thresh = self.config['model']['detection_threshold']
heatmap = heatmap.squeeze()
# print("heatmap: ", heatmap.shape)
pts_nms = getPtsFromHeatmap(heatmap, conf_thresh, nms_dist)
semi_thd_nms_sample = np.zeros_like(heatmap)
semi_thd_nms_sample[pts_nms[1, :].astype(np.int),
pts_nms[0, :].astype(np.int)] = 1
return semi_thd_nms_sample
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))