def getHeatmapCorners(file):
img = cv2.imread(file, 0)
if len(img.shape) > 2:
input = img[np.newaxis]
else:
input = np.tile(img[np.newaxis, :, :, np.newaxis], [1, 1, 1, 3])
prediction = model(
torch.tensor(input.astype(np.float32), device=device).float())
pp = prediction.cpu().data.numpy()
if _subpixel:
pp = np.array(utils.subpixelPrediction(np.copy(pp)))
else:
pp = np.array(utils.argmaxPrediction(np.copy(pp))).astype(
np.float32)
plt.figure(figsize=(15, 15))
for kk in range(1):
plt.subplot(1, 1, 1)
plt.imshow(input[0])
location_list = pp
for kk in range(len(location_list)):
plt.scatter(location_list[kk][0], location_list[kk][1], c='r')
plt.axis('equal')
plt.gca().axis('off')
plt.tight_layout()
plt.savefig('tmpFinal.png')
plt.close('all')
return pp
def getHeatmapCorners(file):
img = cv2.imread(file, 0)
if len(img.shape) > 2:
input = img[np.newaxis]
else:
input = np.tile(img[np.newaxis, :, :, np.newaxis], [1, 1, 1, 3])
prediction = model(
torch.tensor(input.astype(np.float32), device=device).float())
pp = prediction.cpu().data.numpy()
if _subpixel:
pp = np.array(utils.subpixelPrediction(np.copy(pp)))
else:
pp = np.array(utils.argmaxPrediction(np.copy(pp))).astype(np.float32)
return pp
plt.figure(figsize=(15, 15))
for kk in range(4):
plt.subplot(2, 2, kk + 1)
plt.imshow(mesa_images[kk])
plt.imshow(pp[kk], alpha=0.8, cmap='Reds')
plt.axis('equal')
plt.gca().axis('off')
plt.tight_layout()
plt.savefig('plt/' + version + 'benchmark_mesa.png')
plt.close('all')
plt.figure(figsize=(15, 15))
for kk in range(4):
plt.subplot(2, 2, kk + 1)
plt.imshow(mesa_images[kk])
location_list = utils.subpixelPrediction(pp[kk])
for kk in range(len(location_list)):
plt.scatter(location_list[kk][0], location_list[kk][1], c='r')
plt.axis('equal')
plt.gca().axis('off')
plt.tight_layout()
plt.savefig('plt/' + version + 'benchmark_mesa_2.png')
plt.close('all')
# GoPro
print("--- GoPro")
plt.figure(figsize=(15, 15))
for kk in range(4):
pp = model(torch.tensor(gopro_images[kk:kk+1], device=device).float()).cpu().data.numpy()
plt.subplot(1, 1, jj + 1)
plt.imshow(X_data[idx_test + jj])
location_list = utils.argmaxPrediction(np.copy(pp))
for kk in range(len(location_list)):
plt.scatter(location_list[kk][0], location_list[kk][1], c='r')
plt.axis('equal')
plt.gca().axis('off')
plt.tight_layout()
plt.savefig('tmp.png')
plt.close('all')
if _subpixel:
pp = np.array(utils.subpixelPrediction(np.copy(pp)))
else:
pp = np.array(utils.argmaxPrediction(np.copy(pp)))
yy = Y_label[idx_test][Y_label[idx_test][:, 2] > 0.2, :]
distance = np.zeros([pp.shape[0], yy.shape[0]])
for kk in range(pp.shape[0]):
for jj in range(yy.shape[0]):
distance[kk, jj] = np.linalg.norm(pp[kk, :2] - yy[jj, :2])
if distance.shape[0] != 0:
mini = np.min(distance, axis=0)
mini = mini[mini < 3] #non-detection