Python影象讀寫方法對比
訓練視覺相關的神經網路模型時,總是要用到影象的讀寫。方法有很多,比如matplotlib、cv2、PIL等。下面比較幾種讀寫方式,旨在選出一個最快的方式,提升訓練速度。
實驗標準
因為訓練使用的框架是Pytorch,因此讀取的實驗標準如下:
1、讀取解析度都為1920x1080的5張圖片(png格式一張,jpg格式四張)並儲存到陣列。
2、將讀取的陣列轉換為維度順序為CxHxW的Pytorch張量,並儲存到視訊記憶體中(我使用GPU訓練),其中三個通道的順序為RGB。
3、記錄各個方法在以上操作中所耗費的時間。因為png格式的圖片大小差不多是質量有微小差異的jpg格式的10倍,所以資料集通常不會用png來儲存,就不比較這兩種格式的讀取時間差異了。
寫入的實驗標準如下:
1、將5張1920x1080的5張影象對應的Pytorch張量轉換為對應方法可使用的資料型別陣列。
2、以jpg格式儲存五張圖片。
3、記錄各個方法儲存圖片所耗費的時間。
實驗情況
cv2
因為有GPU,所以cv2讀取圖片有兩種方式:
1、先把圖片都讀取為一個numpy陣列,再轉換成儲存在GPU中的pytorch張量。
2、初始化一個儲存在GPU中的pytorch張量,然後將每張圖直接複製進這個張量中。
第一種方式實驗程式碼如下:
import os, torch import cv2 as cv import numpy as np from time import time read_path = 'D:test' write_path = 'D:test\\write\\' # cv2讀取 1 start_t = time() imgs = np.zeros([5, 1080, 1920, 3]) for img, i in zip(os.listdir(read_path), range(5)): img = cv.imread(filename=os.path.join(read_path, img)) imgs[i] = img imgs = torch.tensor(imgs).to('cuda')[...,[2,1,0]].permute([0,3,1,2])/255 print('cv2 讀取時間1:', time() - start_t) # cv2儲存 start_t = time() imgs = (imgs.permute([0,2,3,1])[...,[2,1,0]]*255).cpu().numpy() for i in range(imgs.shape[0]): cv.imwrite(write_path + str(i) + '.jpg', imgs[i]) print('cv2 儲存時間:', time() - start_t)
實驗結果:
cv2 讀取時間1: 0.39693760871887207 cv2 儲存時間: 0.3560612201690674
第二種方式實驗程式碼如下:
import os, torch import cv2 as cv import numpy as np from time import time read_path = 'D:test' write_path = 'D:test\\write\\' # cv2讀取 2 start_t = time() imgs = torch.zeros([5, 1080, 1920, 3], device='cuda') for img, i in zip(os.listdir(read_path), range(5)): img = torch.tensor(cv.imread(filename=os.path.join(read_path, img)), device='cuda') imgs[i] = img imgs = imgs[...,[2,1,0]].permute([0,3,1,2])/255 print('cv2 讀取時間2:', time() - start_t) # cv2儲存 start_t = time() imgs = (imgs.permute([0,2,3,1])[...,[2,1,0]]*255).cpu().numpy() for i in range(imgs.shape[0]): cv.imwrite(write_path + str(i) + '.jpg', imgs[i]) print('cv2 儲存時間:', time() - start_t)
實驗結果:
cv2 讀取時間2: 0.23636841773986816 cv2 儲存時間: 0.3066873550415039
matplotlib
同樣兩種讀取方式,第一種程式碼如下:
import os, torch import numpy as np import matplotlib.pyplot as plt from time import time read_path = 'D:test' write_path = 'D:test\\write\\' # matplotlib 讀取 1 start_t = time() imgs = np.zeros([5, 1080, 1920, 3]) for img, i in zip(os.listdir(read_path), range(5)): img = plt.imread(os.path.join(read_path, img)) imgs[i] = img imgs = torch.tensor(imgs).to('cuda').permute([0,3,1,2])/255 print('matplotlib 讀取時間1:', time() - start_t) # matplotlib 儲存 start_t = time() imgs = (imgs.permute([0,2,3,1])).cpu().numpy() for i in range(imgs.shape[0]): plt.imsave(write_path + str(i) + '.jpg', imgs[i]) print('matplotlib 儲存時間:', time() - start_t)
實驗結果:
matplotlib 讀取時間1: 0.45380306243896484 matplotlib 儲存時間: 0.768944263458252
第二種方式實驗程式碼:
import os, torch import numpy as np import matplotlib.pyplot as plt from time import time read_path = 'D:test' write_path = 'D:test\\write\\' # matplotlib 讀取 2 start_t = time() imgs = torch.zeros([5, 1080, 1920, 3], device='cuda') for img, i in zip(os.listdir(read_path), range(5)): img = torch.tensor(plt.imread(os.path.join(read_path, img)), device='cuda') imgs[i] = img imgs = imgs.permute([0,3,1,2])/255 print('matplotlib 讀取時間2:', time() - start_t) # matplotlib 儲存 start_t = time() imgs = (imgs.permute([0,2,3,1])).cpu().numpy() for i in range(imgs.shape[0]): plt.imsave(write_path + str(i) + '.jpg', imgs[i]) print('matplotlib 儲存時間:', time() - start_t)
實驗結果:
matplotlib 讀取時間2: 0.2044532299041748 matplotlib 儲存時間: 0.4737534523010254
需要注意的是,matplotlib讀取png格式圖片獲取的陣列的數值是在$[0, 1]$範圍內的浮點數,而jpg格式圖片卻是在$[0, 255]$範圍內的整數。所以如果資料集內圖片格式不一致,要注意先轉換為一致再讀取,否則資料集的預處理就麻煩了。
PIL
PIL的讀取與寫入並不能直接使用pytorch張量或numpy陣列,要先轉換為Image型別,所以很麻煩,時間複雜度上肯定也是佔下風的,就不實驗了。
torchvision
torchvision提供了直接從pytorch張量儲存圖片的功能,和上面讀取最快的matplotlib的方法結合,程式碼如下:
import os, torch import matplotlib.pyplot as plt from time import time from torchvision import utils read_path = 'D:test' write_path = 'D:test\\write\\' # matplotlib 讀取 2 start_t = time() imgs = torch.zeros([5, 1080, 1920, 3], device='cuda') for img, i in zip(os.listdir(read_path), range(5)): img = torch.tensor(plt.imread(os.path.join(read_path, img)), device='cuda') imgs[i] = img imgs = imgs.permute([0,3,1,2])/255 print('matplotlib 讀取時間2:', time() - start_t) # torchvision 儲存 start_t = time() for i in range(imgs.shape[0]): utils.save_image(imgs[i], write_path + str(i) + '.jpg') print('torchvision 儲存時間:', time() - start_t)
實驗結果:
matplotlib 讀取時間2: 0.15358829498291016 torchvision 儲存時間: 0.14760661125183105
可以看出這兩個是最快的讀寫方法。另外,要讓圖片的讀寫儘量不影響訓練程序,我們還可以讓這兩個過程與訓練並行。另外,utils.save_image可以將多張圖片拼接成一張來儲存,具體使用方法如下:
utils.save_image(tensor = imgs, # 要儲存的多張圖片張量 shape = [n, C, H, W] fp = 'test.jpg', # 儲存路徑 nrow = 5, # 多圖拼接時,每行所佔的圖片數 padding = 1, # 多圖拼接時,每張圖之間的間距 normalize = True, # 是否進行規範化,通常輸出影象用tanh,所以要用規範化 range = (-1,1)) # 規範化的範圍
&n