如何用Pandas从含像素位置的CSV文件绘制热力图？

问题背景 #

你有一份包含像素坐标（x、y）和对应数值的文本文件，格式如下：

0 0 8.400000e+01
1 0 8.500000e+01
2 0 8.700000e+01
3 0 8.500000e+01
4 0 9.400000e+01
5 0 7.700000e+01
6 0 8.000000e+01
7 0 8.300000e+01
8 0 8.900000e+01
9 0 8.500000e+01
10 0 8.300000e+01

你编写的Python代码运行时报错TypeError: Image data of dtype object cannot be converted to float，代码如下：

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

path_to_csv= "/run/media/test.txt"
df= pd.read_csv(path_to_csv ,sep='\t')
plt.imshow(df,cmap='hot',interpolation='nearest')

plt.show(df)

报错原因及解决步骤 #

1. 修正数据读取逻辑 #

你的数据是空格分隔，不是制表符（\t），且无列名，直接读取会把第一行数据误判为列名，导致数据类型混乱。读取时需指定分隔符和列名：

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

path_to_csv = "/run/media/test.txt"
# 指定列名，用正则匹配任意数量的空格作为分隔符
df = pd.read_csv(path_to_csv, sep=r'\s+', names=['x', 'y', 'value'])
# 强制转换数值列为float类型
df['value'] = df['value'].astype(float)

2. 重塑数据为二维矩阵 #

imshow和seaborn.heatmap都需要二维数组作为输入，当前数据是长格式，需将x、y坐标映射到矩阵的行列位置，填充对应数值：

# 获取x、y的最大索引，确定矩阵尺寸
max_x = df['x'].max()
max_y = df['y'].max()
# 创建空矩阵
heatmap_data = np.zeros((max_y + 1, max_x + 1))
# 遍历数据填充矩阵
for _, row in df.iterrows():
    heatmap_data[row['y'], row['x']] = row['value']

3. 绘制热力图 #

方法一：使用Matplotlib的imshow

plt.imshow(heatmap_data, cmap='hot', interpolation='nearest')
plt.colorbar(label='数值')
plt.xlabel('X坐标')
plt.ylabel('Y坐标')
plt.show()

方法二：使用Seaborn的heatmap

sns.heatmap(heatmap_data, cmap='hot', annot=False, cbar=True)
plt.xlabel('X坐标')
plt.ylabel('Y坐标')
plt.show()

完整可运行代码 #

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

path_to_csv = "/run/media/test.txt"
# 读取并预处理数据
df = pd.read_csv(path_to_csv, sep=r'\s+', names=['x', 'y', 'value'])
df['value'] = df['value'].astype(float)

# 构建热力图矩阵
max_x = df['x'].max()
max_y = df['y'].max()
heatmap_data = np.zeros((max_y + 1, max_x + 1))
for _, row in df.iterrows():
    heatmap_data[row['y'], row['x']] = row['value']

# 绘制Matplotlib热力图
plt.figure(figsize=(8, 4))
plt.imshow(heatmap_data, cmap='hot', interpolation='nearest')
plt.colorbar(label='数值')
plt.xlabel('X坐标')
plt.ylabel('Y坐标')
plt.title('Matplotlib热力图')
plt.show()

# 绘制Seaborn热力图
plt.figure(figsize=(8, 4))
sns.heatmap(heatmap_data, cmap='hot', cbar=True, cbar_kws={'label': '数值'})
plt.xlabel('X坐标')
plt.ylabel('Y坐标')
plt.title('Seaborn热力图')
plt.show()

内容的提问来源于stack exchange，提问作者user979974