【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)