1 PyTorch中語義分割的內置模型

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

2 MaskR-CNN內置模型實現語義分割

2.1 代碼邏輯簡述

將COCO 2017數據集上的預訓練模型dceplabv3_resnet101_coco加載到內存，并使用該模型對圖片進行語義分割。

2.2 代碼實現：MaskR-CNN內置模型實現語義分割

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) # 調用內置模型，并使用預訓練權重進行初始化。 model.eval() # 不然報錯 Expected more than 1 value per channel when training, got input size torch.Size# 在圖片的數據輸入網絡之前，對圖片進行預處理 transform = transforms.Compose([transforms.Resize(256), # 將圖片尺寸調整為256×256transforms.CenterCrop(224), # 中心裁剪成224×224transforms.ToTensor(), # 轉換成張量歸一化到[0,1]transforms.Normalize( # 使用均值，方差標準化mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]) ])def preimg(img): # 定義圖片預處理函數if img.mode == 'RGBA': # 兼容RGBA圖片ch = 4print('ch', ch)a = np.asarray(img)[:, :, :3]img = Image.fromarray(a)return img# 加載要預測的圖片 img = Image.open('./models_2/mask.jpg') # 將圖片輸入模型，進行預測。 # 模型預測的輸出是一個OrderedDict結構。deeplabv3_resnet101模型的圖片輸入尺寸是[224，224]，輸出形狀是[1，21，224，224]，代表20+1(背景）個類別。 plt.imshow(img) plt.axis('off') plt.show() # 顯示加載圖片 im = preimg(img) # 對輸入數據進行維度擴展，成為NCHW inputimg = transform(im).unsqueeze(0)# 顯示用transform轉化后的圖片 tt = np.transpose(inputimg.detach().numpy()[0],(1,2,0)) plt.imshow(tt.astype('uint8')) # 不然報錯：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("輸出結果的形狀：",output['out'].shape) # 去掉批次維度，提取結果。使用argmax函數在每個像素點的21個分類中選出概率值最大的索引，為預測結果。 output = torch.argmax(output['out'].squeeze(), dim=0).detach().cpu().numpy() resultclass = set(list(output.flat)) print("所發現的分類：",resultclass) # 所發現的分類．{0，13，15} # 模型從圖中識別出了兩個類別的內容。索引值13和15分別對應分類名稱“馬”和“人”。def decode_segmap(image,nc=21): # 對圖片中的每個像素點根據其所屬類別進行染色。不同的類別顯示不同的顏色。label_colors = np.array([(0, 0, 0), # 定義每個分類對應的顏色(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): # 根據預測結果進行染色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) # 返回結果rgb = decode_segmap(output) img = Image.fromarray(rgb) plt.axis('off') # 顯示模型的可視化結果 print("快完了") plt.imshow(img) plt.show()

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。