pytorch讀取一張圖像進行分類預測需要注意的問題(opencv、PIL)

讀取圖像一般是兩個庫：opencv和PIL

1、使用opencv讀取圖像

import cv2
image=cv2.imread("/content/drive/My Drive/colab notebooks/image/cat1.jpg")
print(image.shape)

(490, 410, 3)

2、使用PIL讀取圖像

import PIL
image=PIL.Image.open("/content/drive/My Drive/colab notebooks/image/cat1.jpg")
print(image.shape)

這里會報錯：

AttributeError                            Traceback (most recent call last)

<ipython-input-30-807ec7af434b> in <module>()
      1 import PIL
      2 image=PIL.Image.open("/content/drive/My Drive/colab notebooks/image/cat1.jpg")
----> 3 print(image.shape)

AttributeError: 'JpegImageFile' object has no attribute 'shape'

我們要輸出要這么做：

import numpy as np
print(np.array(image).shape)

(490, 410, 3)

需要注意的是：

使用opencv讀取圖像之后是BGR格式的，使用PIL讀取圖像之后是RGB格式的。

3、opencv格式的和PIL格式的之間的轉換

這里參考：https://www.cnblogs.com/enumx/p/12359850.html

（1）opencv格式轉換為PIL格式

import cv2
from PIL import Image
import numpy
 
img = cv2.imread("plane.jpg")
cv2.imshow("OpenCV",img)
image = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
image.show()
cv2.waitKey()

（2）PIL格式轉換為opencv格式

import cv2
from PIL import Image
import numpy
 
image = Image.open("plane.jpg")
image.show()
img = cv2.cvtColor(numpy.asarray(image),cv2.COLOR_RGB2BGR)
cv2.imshow("OpenCV",img)
cv2.waitKey()

4、使用pytorch讀取一張圖片並進行分類預測

需要注意兩個問題：

• 輸入要轉換為：[1，channel，H，W]
• 對輸入的圖像進行數據增強時要求是PIL.Image格式的

import torchvision
import sys
import torch
import torch.nn as nn
from PIL import Image
sys.path.append("/content/drive/My Drive/colab notebooks")
import glob
import numpy as np
import torchvision.transforms as transforms

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=torchvision.models.resnet18(pretrained=False)
model.fc = nn.Linear(model.fc.in_features,4,bias=False)
model.to(device)
model.eval()
save_path="/content/drive/My Drive/colab notebooks/checkpoint/resnet18_best_v2.t7" 
checkpoint = torch.load(save_path)
model.load_state_dict(checkpoint['model'])
print("當前模型准確率為：",checkpoint["epoch_acc"])
images_path="/content/drive/My Drive/colab notebooks/data/dataset/test/four"
transform = transforms.Compose([transforms.Resize((224,224))])
def predict():
  true_labels=[]
  output_labels=[]
  for image in glob.glob(images_path+"/*.png"):
    print(image)
    true_labels.append(0)
    #image=Image.open(image)
    #image=image.resize((224,224))
    image=cv2.imread(image)
    image=cv2.resize(image,(224,224))
    image = Image.fromarray(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
    #print(np.array(image).shape)
    tensor=torch.from_numpy(np.asarray(image)).permute(2,0,1).float()/255.0
    tensor=tensor.reshape((1,3,224,224))
    tensor=tensor.to(device)
    #print(tensor.shape)
    output=model(tensor)
    print(output)
    _, pred = torch.max(output.data,1)
    output_labels.append(pred.item())
  return true_labels,output_labels

true_labels,output_labels=predict()
print("正確的標簽是：")
print(true_labels)
print("預測的標簽是：")
print(output_labels)