python - 将 pandas 数据帧保存到 csv 时如何保留 numpy 数组

Question

我有一个 pandas.DataFrame，其中一列作为图像。该列的每一行都是一个二维 numpy.array 图像。我使用 pandas.DataFrame.to_csv() 将 DataFrame 保存到 csv 文件。但是，当我打开 csv 文件时，该列变为字符串而不是 numpy.array。

如何读取csv文件并保留numpy.array？

Answer 1

要从 csv 文件中读取 numpy.array，您可以为 pandas.read_csv 提供一个转换器函数。

代码:

import ast
import numpy as np
def from_np_array(array_string):
    array_string = ','.join(array_string.replace('[ ', '[').split())
    return np.array(ast.literal_eval(array_string))

测试代码:

import numpy as np
import pandas as pd

image = np.array([[0.1, 0.2], [0.3, 0.4]])
df = pd.DataFrame(
    [['image name1', image],
     ['image name2', image],
     ],
    columns=['names', 'images']).set_index('names')
print(df)
df.to_csv('sample.csv')

df2 = pd.read_csv('sample.csv', converters={'images': from_np_array})
print(df2)

结果:

                               images
names                                
image name1  [[0.1, 0.2], [0.3, 0.4]]
image name2  [[0.1, 0.2], [0.3, 0.4]]

         names                    images
0  image name1  [[0.1, 0.2], [0.3, 0.4]]
1  image name2  [[0.1, 0.2], [0.3, 0.4]]