python读写csv简例

想将数据集中 image 和对应的 label 路径写入一个 csv 文件，避免每次都遍历多个目录。本篇记录 python 读、写 csv 文件的简例，csv 模块用法参考 [1,2]，其中 [1] 有 DictReader、DictWriter 的用例，[2] 有配合 collections.namedtuple 的用法。

分隔符默认是 ,，如果改用 \t（或其它），则 reader、writer 都要改 delimiter 的参数（下文第二例），扩展名可以也相应改成 .tsv，不然 excel 打开会出错。

Code

import csv

# 文件列表
images, labels = [], []
for i in range(10):
    images.append(f"data/image-{i}.png")
    labels.append(f"data/label-{i}.npy")

# 写 csv
with open("test.csv", "w", newline='') as f: # 要对 newline 赋空串
    writer = csv.writer(f, delimiter=',')
    for img, lab in zip(images, labels):
        writer.writerow([img, lab])

# 读 csv
with open("test.csv", "r") as f:
    for line in csv.reader(f):
        print(line)

read & write as dict

类似 [2] 中配合 namedtupe 的用法，不过是用 csv 的接口。跟上节比，csv 文件中多了一行 header。用 DictReader 不用像 [2] 那样手动处理首行 header。

import csv

images, labels = [], []
for i in range(10):
    images.append(f"data/image-{i}.png")
    labels.append(f"data/label-{i}.npy")

with open("test.tsv", "w", newline='') as f:
    field_names = ["image", "label"]
    writer = csv.DictWriter(f, fieldnames=field_names, delimiter='\t')
    writer.writeheader() # 写 header 先
    for il_pair in zip(images, labels):
        row = {fn: fv for fn, fv in zip(field_names, il_pair)} # 以 dict 的形式写行
        writer.writerow(row)

with open("test.tsv", "r") as f:
    for line in csv.DictReader(f, delimiter='\t'): # 不用手动 next 首行 header
        print(line)

References

1. csv — CSV File Reading and Writing
2. python读写csv文件的正确方式