【python】csv文件的读写操作&&将字符串里的数字找出

最近在学习pytorch，有一个猫狗分类的kaggle比赛，预测结果要保存在csv文件中提交。其中格式要求第一列是id（纯数字），第二列是预测为狗的概率。 但我预测出来的id那一列里忘记tensor.item()操作了，以至于输出的不是纯数字，懒得在代码里改了，也抓住这个机会来学一学字符串的处理与csv的读取操作。

这个是预测出来的图，很明显不符合赛方要求的格式

需要将id这一列变成纯数字，如下才是正确的

一、CSV文件的读取操作

使用python的csv库读取

import csv 
file_name = 'test.csv'
with open(file_name, 'r') as f: # 'r'是只读方式读取
    rows = csv.reader(f)
    for row in rows:
        print(row)

以下是输出

['id', 'label']
['tensor(291)', '0.854692101']
['tensor(212)', '0.007425705']
['tensor(4528)', '0.652951479']
['tensor(1513)', '0.58295995']
['tensor(906)', '0.991941571']
['tensor(52)', '0.939000905']
['tensor(11518)', '0.999034882']
['tensor(10489)', '0.998047113']
['tensor(1626)', '0.974285424']
[Finished in 0.1s]

这里提一下open()的读写参数参数：详见：传送门

值得注意的是：

      w模式打开文件，如果文件中有数据，再次写入内容，会把原来数据的覆盖掉

     a模式打开文件，如果文件中有数据，再次写入内容时，不会覆盖原来的数据，而是在其后添加数据

二、提取出字符串中的数字

使用正则表达式提取，需要re库

import csv 
import re 

pattern = re.compile(r'\d+')   # 查找数字

file_name = 'test.csv'
with open(file_name, 'r') as f: # 'r'是只读方式读取
    rows = csv.reader(f)
    for row in rows:
    	num = pattern.findall(row[0]) # 读出第一列中字符串中的数字
    	print(num)
		

以下是输出

[]
['291']
['212']
['4528']
['1513']
['906']
['52']
['11518']
['10489']
['1626']
[Finished in 0.1s]

三、csv文件的写入

代码：

with open('test_write.csv', 'w') as f: # 如果指定的文件不存在，则会自动创建
	writer = csv.writer(f)			
	writer.writerow(['序号', '姓名'])
	

四、将提取出的数字写入到csv文件中

import csv 
import re  

pattern = re.compile(r'\d+')   # 查找数字

# 读取第一列，提取出数字
fname = 'result.csv'
num = []
with open(fname, 'r') as f :
	rows = csv.reader(f)
	for row in rows:
		
		if row[0] == "id" :
			continue
		#print(*pattern.findall(row[0]))	
		num.append([*pattern.findall(row[0]), row[1]])


# 写入到result_new.csv文件里
with open('result_new.csv', 'w', newline='') as f:
	writer = csv.writer(f)			
	writer.writerow(['id', 'label'])
	writer.writerows(num)