1 PyTorch中語義分割的內(nèi)置模型

在torchvision庫下的models\segmentation目錄中，找到segmentation.Py文件。該文件中存放著PyTorch內(nèi)置的語義分割模型。

2 MaskR-CNN內(nèi)置模型實(shí)現(xiàn)語義分割

2.1 代碼邏輯簡述

將COCO 2017數(shù)據(jù)集上的預(yù)訓(xùn)練模型dceplabv3_resnet101_coco加載到內(nèi)存，并使用該模型對圖片進(jìn)行語義分割。

2.2 代碼實(shí)現(xiàn)：MaskR-CNN內(nèi)置模型實(shí)現(xiàn)語義分割

Maskrcnn_resent_Semantic_Segmentation.py

import torch import matplotlib.pyplot as plt from PIL import Image import numpy as np from torchvision import models from torchvision import transforms import os os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"# 獲取模型，如果本地沒有緩存，則下載 model = models.segmentation.deeplabv3_resnet101(pretrained=True) # 調(diào)用內(nèi)置模型，并使用預(yù)訓(xùn)練權(quán)重進(jìn)行初始化。 model.eval() # 不然報(bào)錯 Expected more than 1 value per channel when training, got input size torch.Size# 在圖片的數(shù)據(jù)輸入網(wǎng)絡(luò)之前，對圖片進(jìn)行預(yù)處理 transform = transforms.Compose([transforms.Resize(256), # 將圖片尺寸調(diào)整為256×256transforms.CenterCrop(224), # 中心裁剪成224×224transforms.ToTensor(), # 轉(zhuǎn)換成張量歸一化到[0,1]transforms.Normalize( # 使用均值，方差標(biāo)準(zhǔn)化mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]) ])def preimg(img): # 定義圖片預(yù)處理函數(shù)if img.mode == 'RGBA': # 兼容RGBA圖片ch = 4print('ch', ch)a = np.asarray(img)[:, :, :3]img = Image.fromarray(a)return img# 加載要預(yù)測的圖片 img = Image.open('./models_2/mask.jpg') # 將圖片輸入模型，進(jìn)行預(yù)測。 # 模型預(yù)測的輸出是一個OrderedDict結(jié)構(gòu)。deeplabv3_resnet101模型的圖片輸入尺寸是[224，224]，輸出形狀是[1，21，224，224]，代表20+1(背景）個類別。 plt.imshow(img) plt.axis('off') plt.show() # 顯示加載圖片 im = preimg(img) # 對輸入數(shù)據(jù)進(jìn)行維度擴(kuò)展，成為NCHW inputimg = transform(im).unsqueeze(0)# 顯示用transform轉(zhuǎn)化后的圖片 tt = np.transpose(inputimg.detach().numpy()[0],(1,2,0)) plt.imshow(tt.astype('uint8')) # 不然報(bào)錯：Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers) plt.show()output = model(inputimg) # 將圖片輸入模型 print("輸出結(jié)果的形狀：",output['out'].shape) # 去掉批次維度，提取結(jié)果。使用argmax函數(shù)在每個像素點(diǎn)的21個分類中選出概率值最大的索引，為預(yù)測結(jié)果。 output = torch.argmax(output['out'].squeeze(), dim=0).detach().cpu().numpy() resultclass = set(list(output.flat)) print("所發(fā)現(xiàn)的分類：",resultclass) # 所發(fā)現(xiàn)的分類．{0，13，15} # 模型從圖中識別出了兩個類別的內(nèi)容。索引值13和15分別對應(yīng)分類名稱“馬”和“人”。def decode_segmap(image,nc=21): # 對圖片中的每個像素點(diǎn)根據(jù)其所屬類別進(jìn)行染色。不同的類別顯示不同的顏色。label_colors = np.array([(0, 0, 0), # 定義每個分類對應(yīng)的顏色(128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),(0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),(192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),(0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])r = np.zeros_like(image).astype(np.uint8) # 初始化RGBg = np.zeros_like(image).astype(np.uint8)b = np.zeros_like(image).astype(np.uint8)for l in range(0, nc): # 根據(jù)預(yù)測結(jié)果進(jìn)行染色idx = image == lprint("idx：",idx)r[idx] = label_colors[l, 0]g[idx] = label_colors[l, 1]b[idx] = label_colors[l, 2]return np.stack([r, g, b], axis=2) # 返回結(jié)果rgb = decode_segmap(output) img = Image.fromarray(rgb) plt.axis('off') # 顯示模型的可視化結(jié)果 print("快完了") plt.imshow(img) plt.show()

總結(jié)

以上是生活随笔為你收集整理的【Pytorch神经网络实战案例】27 MaskR-CNN内置模型实现语义分割的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。