輕松學(xué)Pytorch-使用ResNet50實現(xiàn)圖像分類

Hello大家好，這篇文章給大家詳細(xì)介紹一下pytorch中最重要的組件torchvision，它包含了常見的數(shù)據(jù)集、模型架構(gòu)與預(yù)訓(xùn)練模型權(quán)重文件、常見圖像變換、計算機(jī)視覺任務(wù)訓(xùn)練。可以是說是pytorch中非常有用的模型遷移學(xué)習(xí)神器。本文將會介紹如何使用torchvison的預(yù)訓(xùn)練模型ResNet50實現(xiàn)圖像分類。

Torchvision.models包里面包含了常見的各種基礎(chǔ)模型架構(gòu)，主要包括：

AlexNet
VGG
ResNet
SqueezeNet
DenseNet
Inception v3
GoogLeNet
ShuffleNet v2
MobileNet v2
ResNeXt
Wide ResNet
MNASNet

這里我選擇了ResNet50，基于ImageNet訓(xùn)練的基礎(chǔ)網(wǎng)絡(luò)來實現(xiàn)圖像分類， 網(wǎng)絡(luò)模型下載與加載如下：

model = torchvision.models.resnet50(pretrained=True).eval().cuda()tf = transforms.Compose([            transforms.Resize(256),            transforms.CenterCrop(224),            transforms.ToTensor(),            transforms.Normalize(            mean=[0.485, 0.456, 0.406],            std=[0.229, 0.224, 0.225]        )])

這里首先需要加載ImageNet的分類標(biāo)簽，目的是最后顯示分類的文本標(biāo)簽時候使用。然后對輸入圖像完成預(yù)處理，使用ResNet50模型實現(xiàn)分類預(yù)測，對預(yù)測結(jié)果解析之后，顯示標(biāo)簽文本，完整的代碼演示如下：

 1with open('imagenet_classes.txt') as f:
 2    labels = [line.strip() for line in f.readlines()]
 3
 4src = cv.imread("D:/images/space_shuttle.jpg") # aeroplane.jpg
 5image = cv.resize(src, (224, 224))
 6image = np.float32(image) / 255.0
 7image[:,:,] -= (np.float32(0.485), np.float32(0.456), np.float32(0.406))
 8image[:,:,] /= (np.float32(0.229), np.float32(0.224), np.float32(0.225))
 9image = image.transpose((2, 0, 1))
10input_x = torch.from_numpy(image).unsqueeze(0)
11print(input_x.size())
12pred = model(input_x.cuda())
13pred_index = torch.argmax(pred, 1).cpu().detach().numpy()
14print(pred_index)
15print("current predict class name : %s"%labels[pred_index[0]])
16cv.putText(src, labels[pred_index[0]], (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
17cv.imshow("input", src)
18cv.waitKey(0)
19cv.destroyAllWindows()

運行結(jié)果如下：

在torchvision中的模型基本上都可以轉(zhuǎn)換為ONNX格式，而且被OpenCV DNN模塊所支持，所以，很方便的可以對torchvision自帶的模型轉(zhuǎn)為ONNX，實現(xiàn)OpenCV DNN的調(diào)用，首先轉(zhuǎn)為ONNX模型，直接使用torch.onnx.export即可轉(zhuǎn)換(還不知道怎么轉(zhuǎn)，快點看前面的例子)。轉(zhuǎn)換之后使用OpenCV DNN調(diào)用的代碼如下：

 1with open('imagenet_classes.txt') as f:
 2    labels = [line.strip() for line in f.readlines()]
 3net = cv.dnn.readNetFromONNX("resnet.onnx")
 4src = cv.imread("D:/images/messi.jpg")  # aeroplane.jpg
 5image = cv.resize(src, (224, 224))
 6image = np.float32(image) / 255.0
 7image[:, :, ] -= (np.float32(0.485), np.float32(0.456), np.float32(0.406))
 8image[:, :, ] /= (np.float32(0.229), np.float32(0.224), np.float32(0.225))
 9blob = cv.dnn.blobFromImage(image, 1.0, (224, 224), (0, 0, 0), False)
10net.setInput(blob)
11probs = net.forward()
12index = np.argmax(probs)
13cv.putText(src, labels[index], (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
14cv.imshow("input", src)
15cv.waitKey(0)
16cv.destroyAllWindows()

 運行結(jié)果見上圖，這里就不再貼了。

好消息！
小白學(xué)視覺知識星球
開始面向外開放啦??????



下載1：OpenCV-Contrib擴(kuò)展模塊中文版教程
在「小白學(xué)視覺」公眾號后臺回復(fù)：擴(kuò)展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴(kuò)展模塊教程中文版，涵蓋擴(kuò)展模塊安裝、SFM算法、立體視覺、目標(biāo)跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。

下載2：Python視覺實戰(zhàn)項目52講
在「小白學(xué)視覺」公眾號后臺回復(fù)：Python視覺實戰(zhàn)項目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內(nèi)容提取、面部識別等31個視覺實戰(zhàn)項目，助力快速學(xué)校計算機(jī)視覺。

下載3：OpenCV實戰(zhàn)項目20講
在「小白學(xué)視覺」公眾號后臺回復(fù)：OpenCV實戰(zhàn)項目20講，即可下載含有20個基于OpenCV實現(xiàn)20個實戰(zhàn)項目，實現(xiàn)OpenCV學(xué)習(xí)進(jìn)階。

交流群

歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計算攝影、檢測、分割、識別、醫(yī)學(xué)影像、GAN、算法競賽等微信群（以后會逐漸細(xì)分），請掃描下面微信號加群，備注：”昵稱+學(xué)校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進(jìn)入相關(guān)微信群。請勿在群內(nèi)發(fā)送廣告，否則會請出群，謝謝理解~