根据pytorch预测生成的热力图在原图画点

在使用pytorch框架进行开发时，有时需要根据预测结果在原图画点，下面是实现的方法：

具体代码的含义将通过注释的形式进行标注：

img_r = cv2.imread("test21/room/1.jpg") # 读入图片
img1 = cv2.imread("test21/room/1.jpg", 0)
img = img1.astype(np.float32, copy=False)
ht = img.shape[0]
wd = img.shape[1]
ht_1 = int((ht/4)*4)
wd_1 = int((wd/4)*4)
img = cv2.resize(img,(wd_1, ht_1))# 重置图片大小
img = img.reshape((1, 1, img.shape[0], img.shape[1]))
heat_map = net.forward(img)# 根据生成的模型进行预测
heat_map = torch.squeeze(heat_map)
heat_map = heat_map.data.numpy()# 将生成的结果转化为numpy数组
[height, width] = heat_map.shape
dst = cv2.resize(img_r, (width, height), interpolation=cv2.INTER_LINEAR)
print(height)
print(width)
for i in range(width - 1):# 遍历生成的numpy数组
    for j in range(height-1):
       if heat_map[j, i] > 0:
           dst = cv2.circle(dst, (i, j), 1, (0, 0, 255), -1) # 根据坐标在原图画点
cv2.namedWindow("result")
cv2.imshow("result", dst)

cv2.waitKey(0)

 

欢迎关注我的公众号：

编程技术与生活（ID：hw_cchang）