基于深度學習的語意分割初探FCN以及pytorch代碼實作
FCN論文
論文地址：https://arxiv.org/abs/1411.4038
FCN是基于深度學習方法的第一篇關于語意分割的開山之作，雖然這篇文章的分割結果現在看起來并不是目前最好的，但其意義還是非常重要的，其中跳躍鏈接、end-to-end、遷移學習、反卷積實作上采樣也是FCN論文中的核心思想，
FCN論文整體結構
應用
無人車、地理資訊系統、醫療影像、機器人，由于目前想在機器人上搭建視覺系統，想結合語意分割這種像素級預測的思想，是否可以與檢測任務中的方式做一個結合，例如Mask-RCNN將實體分割與目標檢測很好的融合為一體，
pytorch實作FCN_8x
由于讓我們的代碼更加易于理解以及更好的更正，代碼中所有引數以及變數名稱均使用我們的母語，
簡單使用Camvid資料集做一個室外分割的例子，
Dataset構建
import torch
import os
from PIL import Image
import pandas as pd
import numpy as np
import torchvision.transforms.functional as F
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import cfg
class 標簽處理:
def __init__(self, 標簽所對應類別檔案的路徑):
self.像素類別圖 = self.讀取類別所對應的像素值(標簽所對應類別檔案的路徑)
self.標簽哈希表 = self.編碼標簽像素值(self.像素類別圖)
@staticmethod
def 讀取類別所對應的像素值(標簽所對應類別檔案的路徑):
標簽像素值 = pd.read_csv(標簽所對應類別檔案的路徑, sep=',')
像素類別圖 = []
#標簽像素值.index # 回傳像素值所對應類別的索引 0-12
for i in range(len(標簽像素值.index)):
按行讀取每一個類別所對應的像素值 = 標簽像素值.iloc[i]
類別所對應的RGB像素值 = [按行讀取每一個類別所對應的像素值['r'], 按行讀取每一個類別所對應的像素值['g'], 按行讀取每一個類別所對應的像素值['b']]
像素類別圖.append(類別所對應的RGB像素值)
# 類別名稱 = 標簽像素值['name'].values
# 類別數量 = len(類別名稱)
return 像素類別圖
@staticmethod
def 編碼標簽像素值(像素類別圖):
# 哈希表(為了形成1對1或1對多的映射關系，加快查找的效率) 一個標簽對應一個顏色 將像素類別圖中的每一個像素映射到它所表示的類別
# 希函式 像素類別圖([0]*256+像素類別圖[1])*256+像素類別圖[2]
# 哈希映射 像素類別圖2lbl(希函式) = 所對應的類別
# 哈希表 像素類別圖2lbl
# eg: 一個像素點P（128， 64， 128） 通過編碼函式（P[0]*256+P[1]）*256+P[2] 轉成 整數（8405120）
# 將該數作為像素點P在哈希表中的索引：像素類別圖轉成哈希表（8405120） 去查詢像素點P所對應的類別P
像素類別圖轉成哈希表 = np.zeros(256 ** 3)
for 類別索引, 類別所對應RGB像素值 in enumerate(像素類別圖):
像素類別圖轉成哈希表[(類別所對應RGB像素值[0]*256 + 類別所對應RGB像素值[1]) * 256 + 類別所對應RGB像素值[2]] = 類別索引
return 像素類別圖轉成哈希表
def 編碼標簽影像(self, 影像):
# rgb -> index -> identity
資料 = np.array(影像, dtype='int32')
哈希函式值 = (資料[:, :, 0] * 256 + 資料[:, :, 1]) * 256 + 資料[:, :, 2]
return np.array(self.標簽哈希表[哈希函式值], dtype='int64')
class 資料集(Dataset):
def __init__(self, 影像和標簽路徑=[], 裁剪=None):
if len(影像和標簽路徑) != 2:
raise Exception('需同時輸入影像和標簽的路徑')
self.影像路徑 = 影像和標簽路徑[0]
self.標簽路徑 = 影像和標簽路徑[1]
self.讀取路徑中的圖片 = self.讀取檔案夾(self.影像路徑)
self.讀取路徑中的標簽 = self.讀取檔案夾(self.標簽路徑)
self.裁剪尺寸 = 裁剪
def __getitem__(self, 索引):
單張影像 = self.讀取路徑中的圖片[索引]
單個標簽 = self.讀取路徑中的標簽[索引]
單張影像 = Image.open(單張影像)
單個標簽 = Image.open(單個標簽).convert('RGB')
單張影像, 單個標簽 = self.中心裁剪(單張影像, 單個標簽, self.裁剪尺寸)
單張影像, 單個標簽 = self.影像標簽轉換(單張影像, 單個標簽)
影像標簽組合成字典 = {'影像': 單張影像, '標簽': 單個標簽}
return 影像標簽組合成字典
def __len__(self):
return len(self.讀取路徑中的圖片)
def 讀取檔案夾(self, 路徑):
檔案夾串列 = os.listdir(路徑)
拼接影像完整路徑 = [os.path.join(路徑, 圖片) for 圖片 in 檔案夾串列]
拼接影像完整路徑.sort()
return 拼接影像完整路徑
def 中心裁剪(self, 影像, 標簽, 裁剪尺寸):
影像 = F.center_crop(影像, 裁剪尺寸)
標簽 = F.center_crop(標簽, 裁剪尺寸)
return 影像, 標簽
def 影像標簽轉換(self, 影像, 標簽):
標簽 = np.array(標簽)
標簽 = Image.fromarray(標簽.astype('uint8'))
影像轉Tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
影像 = 影像轉Tensor(影像)
# 原圖不需要編碼 標簽需要編碼
標簽 = 標簽處理實體化.編碼標簽影像(標簽)
標簽 = torch.from_numpy(標簽)
return 影像, 標簽
標簽處理實體化 = 標簽處理(cfg.類別檔案路徑)
FCN模型搭建
import torch
import torch.nn as nn
from torchvision import models
import torch.nn.functional as F
import numpy as np
from Bilinear_init_deconv import 雙線性插值初始化卷積核
VGG特征提取網路 = models.vgg16_bn(pretrained=True)
class 全卷積網路(nn.Module):
def __init__(self, 類別個數):
super(全卷積網路, self).__init__()
self.特征提取網路中第一個下采樣 = VGG特征提取網路.features[:7] # 64
self.特征提取網路中第二個下采樣 = VGG特征提取網路.features[7:14] # 128
self.特征提取網路中第三個下采樣 = VGG特征提取網路.features[14:24] # 256
self.特征提取網路中第四個下采樣 = VGG特征提取網路.features[24:34] # 512
self.特征提取網路中第五個下采樣 = VGG特征提取網路.features[34:] # 512
# self.跨度_32的上采樣預測圖 = nn.Conv2d(512, 類別個數, 1) # 32
# self.跨度_16的采樣預測圖 = nn.Conv2d(512, 類別個數, 1) # 16
# self.跨度_8的上采樣預測圖 = nn.Conv2d(128, 類別個數, 1) # 8
self.過渡卷積512 = nn.Conv2d(512, 256, 1)
self.過渡卷積256 = nn.Conv2d(256, 類別個數, 1)
self.上采樣_8X = nn.ConvTranspose2d(類別個數, 類別個數, 16, 8, 4, bias=False)
self.上采樣_8X.weight.data = 雙線性插值初始化卷積核(類別個數, 類別個數, 16)
self.上采樣_2X_512 = nn.ConvTranspose2d(512, 512, 4, 2, 1, bias=False)
self.上采樣_2X_512.weight.data = 雙線性插值初始化卷積核(512, 512, 4)
self.上采樣_2X_256 = nn.ConvTranspose2d(256, 256, 4, 2, 1, bias=False)
self.上采樣_2X_256.weight.data = 雙線性插值初始化卷積核(256, 256, 4)
def forward(self, x):
第一層特征提取 = self.特征提取網路中第一個下采樣(x)
第二層特征提取 = self.特征提取網路中第二個下采樣(第一層特征提取)
第三層特征提取 = self.特征提取網路中第三個下采樣(第二層特征提取)
第四層特征提取 = self.特征提取網路中第四個下采樣(第三層特征提取)
第五層特征提取 = self.特征提取網路中第五個下采樣(第四層特征提取)
第五層特征提取_2倍還原 = self.上采樣_2X_512(第五層特征提取)
第五層與第四層進行特征圖融合 = 第四層特征提取 + 第五層特征提取_2倍還原
融合后的影像轉換通道數 = self.過渡卷積512(第五層與第四層進行特征圖融合)
第四層與第五層融合后的特征_2倍還原 = self.上采樣_2X_256(融合后的影像轉換通道數)
與第三層特征圖進行融合 = 第三層特征提取 + 第四層與第五層融合后的特征_2倍還原
轉換成類別個數的通道數 = self.過渡卷積256(與第三層特征圖進行融合)
還原原圖大小_8X = self.上采樣_8X(轉換成類別個數的通道數)
return 還原原圖大小_8X
FCN論文中使用了雙線性插值初始化反卷積核
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import cv2
def 雙線性插值(原圖, 目標尺寸):
目標影像的高, 目標影像的寬 = 目標尺寸
原圖的高, 原圖的寬 = 原圖.shape[:2]
if 原圖的高 == 目標影像的高 and 原圖的寬 == 目標影像的寬:
return 原圖.copy()
原圖與目標影像寬的縮放比例 = float(原圖的寬) / 目標影像的寬
原圖與目標影像高的縮放比例 = float(原圖的高) / 目標影像的高
生成目標影像尺寸相同的空白圖 = np.zeros((目標影像的高, 目標影像的寬, 3), dtype=np.uint8)
for RGB in range(3):
for 目標影像高方向 in range(目標影像的高):
for 目標影像寬方向 in range(目標影像的寬):
# src_x + 0.5 = (dst_x + 0.5) * scale_x 0.5為一個像素默認1*1 其中心像素坐標+0.5的位置
目標影像寬方向的像素在原圖上的坐標 = (目標影像寬方向 + 0.5) * 原圖與目標影像寬的縮放比例 - 0.5
目標影像高方向的像素在原圖上的坐標 = (目標影像高方向 + 0.5) * 原圖與目標影像高的縮放比例 - 0.5
原圖上第一個近鄰點 = int(np.floor(目標影像寬方向的像素在原圖上的坐標))
原圖上第二個近鄰點 = int(np.floor(目標影像高方向的像素在原圖上的坐標))
原圖上第三個近鄰點 = min(原圖上第一個近鄰點 + 1, 原圖的寬 - 1)
原圖上第四個近鄰點 = min(原圖上第二個近鄰點 + 1, 原圖的高 - 1)
比例1 = (原圖上第三個近鄰點 - 目標影像寬方向的像素在原圖上的坐標) * 原圖[原圖上第二個近鄰點, 原圖上第一個近鄰點, RGB] + (目標影像寬方向的像素在原圖上的坐標 - 原圖上第一個近鄰點) * 原圖[原圖上第二個近鄰點, 原圖上第三個近鄰點, RGB]
比例2 = (原圖上第三個近鄰點 - 目標影像寬方向的像素在原圖上的坐標) * 原圖[原圖上第四個近鄰點, 原圖上第一個近鄰點, RGB] + (目標影像寬方向的像素在原圖上的坐標 - 原圖上第一個近鄰點) * 原圖[原圖上第四個近鄰點, 原圖上第三個近鄰點, RGB]
生成目標影像尺寸相同的空白圖[目標影像高方向, 目標影像寬方向, RGB] = int((原圖上第四個近鄰點 - 原圖上第二個近鄰點) * 比例1 + (目標影像高方向的像素在原圖上的坐標 - 原圖上第二個近鄰點) * 比例2)
return 生成目標影像尺寸相同的空白圖
def 雙線性插值初始化卷積核(輸入通道, 輸出通道, 卷積核大小):
因子 = (卷積核大小 + 1) // 2
if 卷積核大小 % 2 == 1:
中心 = 因子 - 1
else:
中心 = 因子 - 0.5
畫網格 = np.ogrid[:卷積核大小, :卷積核大小]
初始化 = (1 - abs(畫網格[0] - 中心) / 因子) * (1 - abs(畫網格[1] - 中心) / 因子)
權重 = np.zeros((輸入通道, 輸出通道, 卷積核大小, 卷積核大小), dtype='float32')
權重[range(輸入通道), range(輸出通道), :, :] = 初始化
return torch.from_numpy(權重)
if __name__ == '__main__':
img = cv2.imread('FCN_model.png')
img_out = 雙線性插值(img, (1000, 1000))
cv2.imshow('src', img)
cv2.imshow('dst', img_out)
cv2.waitKey(0)
print(img.shape)
print(img_out.shape)
訓練
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import evalution_segmentation
import cfg
from dataset import 資料集
from build_FCN_model import 全卷積網路
from datetime import datetime
計算單元 = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
訓練資料集實體化 = 資料集([cfg.訓練資料集, cfg.訓練標簽資料集], (352, 480))
驗證資料集實體化 = 資料集([cfg.驗證資料集, cfg.驗證標簽資料集], (352, 480))
訓練資料 = DataLoader(訓練資料集實體化, batch_size=16, shuffle=True, num_workers=0)
驗證資料 = DataLoader(驗證資料集實體化, batch_size=8, shuffle=True, num_workers=0)
模型實體化 = 全卷積網路(類別個數=12)
模型放到GPU = 模型實體化.to(計算單元)
損失函式 = nn.NLLLoss().to(計算單元) # 交叉熵沒有本質區別 只是沒有封裝softmax
優化器 = optim.Adam(模型放到GPU.parameters(), lr=1e-4) # 2D Adam rgb-D SGD
def 訓練(模型):
最優權重 = [0]
網路狀態 = 模型.train()
for 訓練輪次 in range(cfg.回圈資料集的總次數):
print('訓練次數[{} / {}]'.format(訓練輪次 + 1, cfg.回圈資料集的總次數))
if 訓練輪次 % 50 == 0 and 訓練輪次 != 0:
for 學習率 in 優化器.param_groups:
學習率['lr'] *= 0.5
訓練損失 = 0
訓練準確率 = 0
訓練miou = 0
訓練分類的準確率 = 0
for 索引, 影像標簽資料字典 in enumerate(訓練資料):
訓練影像資料 = Variable(影像標簽資料字典['影像'].to(計算單元))
訓練影像標簽 = Variable(影像標簽資料字典['標簽'].to(計算單元))
預測圖獲取 = 網路狀態(訓練影像資料)
預測圖獲取 = F.log_softmax(預測圖獲取, dim=1)
損失 = 損失函式(預測圖獲取, 訓練影像標簽) # 每一次迭代的loss
優化器.zero_grad()
損失.backward()
優化器.step()
訓練損失 += 損失.item() # 對于一個epoch總的loss
預測結果中取最大值 = 預測圖獲取.max(dim=1)[1].data.cpu().numpy() # max 回傳兩個值 1、最大值本身 2、最大值的索引
預測結果中取最大值 = [序號 for 序號 in 預測結果中取最大值]
真實標簽資料 = 訓練影像標簽.data.cpu().numpy()
真實標簽資料 = [序號 for 序號 in 真實標簽資料]
混淆矩陣 = evalution_segmentation.驗證語意分割指標(預測結果中取最大值, 真實標簽資料)
訓練準確率 += 混淆矩陣['平均分類精度']
訓練miou += 混淆矩陣['miou']
訓練分類的準確率 += 混淆矩陣['分類精度']
print('迭代到第[{} / {}]個資料, 損失為 {:.8f}'.format(索引 + 1, len(訓練資料), 損失.item()))
每一個大回圈下的指標描述 = '訓練準確率: {:.5f} 訓練miou: {:.5f} 訓練類別的準確率: {:}'.format(訓練準確率 / len(訓練資料), 訓練miou / len(訓練資料), 訓練分類的準確率 / len(訓練資料))
print(每一個大回圈下的指標描述)
if max(最優權重) <= 訓練miou / len(訓練資料):
保存權重路徑 = 'E:/FCN_8X/weights/'
for name in os.listdir(保存權重路徑):
os.remove(os.path.join(保存權重路徑, name))
最優權重.append(訓練miou / len(訓練資料))
torch.save(網路狀態.state_dict(), 保存權重路徑 + '{}.pth'.format(訓練輪次))
驗證(模型)
def 驗證(模型):
網路狀態 = 模型.eval()
驗證損失 = 0
驗證準確度 = 0
驗證miou = 0
驗證分類的準確率 = 0
預測初始時間 = datetime.now()
for 序號, 影像標簽資料 in enumerate(驗證資料):
驗證影像資料 = Variable(影像標簽資料['影像'].to(計算單元))
驗證影像標簽 = Variable(影像標簽資料['標簽'].to(計算單元))
預測圖輸出 = 網路狀態(驗證影像資料)
預測圖輸出 = F.log_softmax(預測圖輸出, dim=1)
損失 = 損失函式(預測圖輸出, 驗證影像標簽)
驗證損失 = 損失.item() + 驗證損失
預測結果中取最大值 = 預測圖輸出.max(dim=1)[1].data.cpu().numpy() # max 回傳兩個值 1、最大值本身 2、最大值的索引
預測結果中取最大值 = [序號 for 序號 in 預測結果中取最大值]
真實標簽資料 = 驗證影像標簽.data.cpu().numpy()
真實標簽資料 = [序號 for 序號 in 真實標簽資料]
驗證混淆矩陣 = evalution_segmentation.驗證語意分割指標(預測結果中取最大值, 真實標簽資料)
驗證準確度 = 驗證混淆矩陣['平均分類精度'] + 驗證準確度
驗證miou += 驗證混淆矩陣['miou']
驗證分類的準確率 += 驗證混淆矩陣['分類精度']
當前時間 = datetime.now()
小時, 分秒 = divmod((當前時間 - 預測初始時間).seconds, 3600)
分鐘, 秒 = divmod(分秒, 60)
驗證用時 = '驗證所用時間為： {:.0f}:{:.0f}:{:.0f}'.format(小時, 分鐘, 秒)
驗證的指標描述 = ('驗證的損失: {:.5f} 驗證的準確度: {:.5f} 驗證的miou: {:.5f} 驗證類精度: {:}').format(驗證損失 / len(訓練資料), 驗證準確度 / len(驗證資料), 驗證miou / len(驗證資料), 驗證分類的準確率 / len(驗證資料))
print(驗證的指標描述)
print(驗證用時)
if __name__ == '__main__':
訓練(模型放到GPU)
測驗
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import evalution_segmentation
import cfg
from dataset import 資料集
from build_FCN_model import 全卷積網路
from datetime import datetime
計算單元 = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
miou_list = [0]
權重路徑 = ''
測驗資料實體化 = 資料集([], (352, 480))
測驗資料 = DataLoader(測驗資料實體化, batch_size=4, shuffle=True, num_workers=0)
網路 = 全卷積網路(類別個數=12)
網路.eval()
網路.to(計算單元)
網路.load_state_dict(torch.load(權重路徑))
測驗的準確度 = 0
測驗的miou = 0
測驗的分類精度 = 0
測驗像素準確度 = 0
錯誤 = 0
for 索引, 影像標簽資料字典 in enumerate(測驗資料):
測驗影像資料 = Variable(影像標簽資料字典['影像'].to(計算單元))
測驗影像標簽 = Variable(影像標簽資料字典['標簽'].to(計算單元))
預測 = 網路(測驗影像資料)
預測 = F.log_softmax(預測, dim=1)
預測結果中取最大值 = 預測.max(dim=1)[1].data.cpu().numpy() # max 回傳兩個值 1、最大值本身 2、最大值的索引
預測結果中取最大值 = [序號 for 序號 in 預測結果中取最大值]
真實標簽資料 = 測驗影像標簽.data.cpu().numpy()
真實標簽資料 = [序號 for 序號 in 真實標簽資料]
測驗混淆矩陣 = evalution_segmentation.驗證語意分割指標(預測結果中取最大值, 真實標簽資料)
測驗的準確度 += 測驗混淆矩陣['平均分類精度']
測驗的miou += 測驗混淆矩陣['miou']
測驗像素準確度 += 測驗混淆矩陣['像素準確度']
if len(測驗混淆矩陣['分類精度']) < 12:
測驗混淆矩陣['分類精度'] = 0
測驗的分類精度 += 測驗混淆矩陣['分類精度']
錯誤 += 1
else:
測驗的分類精度 += 測驗混淆矩陣['分類精度']
print(測驗混淆矩陣['分類精度'], '=============', 索引)
完整回圈一次的指標 = ('測驗精度: {:.5f}, 測驗miou: {:.5f}, 測驗像素準確度: {:.5f}, 測驗分類精度: {}'.format(測驗的準確度 / (len(測驗資料) - 錯誤),
測驗的miou / (len(測驗資料) - 錯誤), 測驗像素準確度 / (len(測驗資料) - 錯誤),
測驗的分類精度 / (len(測驗資料) - 錯誤)))
if 測驗的miou / (len(測驗資料) - 錯誤) > max(miou_list):
miou_list.append(測驗的miou / (len(測驗資料) - 錯誤))
print(完整回圈一次的指標 + '=============')
評價指標
import numpy as np
import six
def 計算混淆矩陣(預測值, 真實標簽):
預測值 = iter(預測值)
真實標簽 = iter(真實標簽)
類別數量 = 12
混淆矩陣 = np.zeros((類別數量, 類別數量), dtype=np.int64)
for 單個預測值, 單個真實標簽 in six.moves.zip(預測值, 真實標簽):
if 單個預測值.ndim != 2 or 單個真實標簽.ndim != 2:
raise ValueError('預測值或標簽必須為2維')
if 單個預測值.shape != 單個真實標簽.shape:
raise ValueError('預測值和標簽的尺寸必須相同')
預測值變換成一維向量 = 單個預測值.flatten()
標簽變換成一維向量 = 單個真實標簽.flatten()
預測和真實值中最大的類別索引 = np.max((預測值變換成一維向量, 標簽變換成一維向量))
if 預測和真實值中最大的類別索引 >= 類別數量:
擴大混淆矩陣 = np.zeros((預測和真實值中最大的類別索引 + 1, 預測和真實值中最大的類別索引 + 1), dtype=np.int64)
擴大混淆矩陣[0:類別數量, 0:類別數量] = 混淆矩陣
類別數量 = 預測和真實值中最大的類別索引 + 1
混淆矩陣 = 擴大混淆矩陣
掩碼 = 單個真實標簽 >= 0
混淆矩陣 += np.bincount(類別數量 * 單個真實標簽[掩碼].astype(int) + 單個預測值[掩碼], minlength=類別數量 ** 2).reshape((類別數量, 類別數量)) # N*L+P
for iter_ in (預測值, 真實標簽):
# This code assumes any iterator does not contain None as its items.
if next(iter_, None) is not None:
raise ValueError('Length of input iterables need to be same')
return 混淆矩陣
def 計算語意分割的iou(混淆矩陣):
iou并集 = (混淆矩陣.sum(axis=1) + 混淆矩陣.sum(axis=0) - np.diag(混淆矩陣)) # 0列 1行
iou = np.diag(混淆矩陣) / iou并集
return iou[:-1] # 最末尾為背景 舍棄
def 驗證語意分割指標(預測值, 真實標簽):
混淆矩陣 = 計算混淆矩陣(預測值, 真實標簽)
iou = 計算語意分割的iou(混淆矩陣)
像素精度 = np.diag(混淆矩陣).sum() / 混淆矩陣.sum()
類精度 = np.diag(混淆矩陣) / (np.sum(混淆矩陣, axis=1) + 1e10)
return {
'iou': iou,
'miou': np.nanmean(iou),
'像素準確度': 像素精度,
'分類精度': 類精度,
'平均分類精度': np.nanmean(類精度[:-1])
}
預測影像
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import pandas as pd
import numpy as np
import evalution_segmentation
import cfg
from dataset import 資料集
from build_FCN_model import 全卷積網路
import datetime
from PIL import Image
計算單元 = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
權重路徑 = 'E:/FCN_8X/weights/194.pth'
標簽檔案路徑 = 'E:/Camvid-FCN/Datasets/CamVid/class_dict.csv'
測驗資料實體化 = 資料集([cfg.測驗資料集, cfg.測驗標簽資料集], (352, 480))
測驗資料 = DataLoader(測驗資料實體化, batch_size=4, shuffle=True, num_workers=0)
網路 = 全卷積網路(類別個數=12).to(計算單元)
網路.load_state_dict(torch.load(權重路徑))
網路.eval()
預測標簽的顏色 = pd.read_csv(標簽檔案路徑, sep=',')
標簽所對應的類別 = 預測標簽的顏色['name'].values
標簽類別的數量 = len(標簽所對應的類別)
像素類別圖 = []
for 類別 in range(標簽類別的數量):
按行讀取每一個類別所對應的像素值 = 預測標簽的顏色.iloc[類別]
類別所對應的RGB像素值 = [按行讀取每一個類別所對應的像素值['r'], 按行讀取每一個類別所對應的像素值['g'], 按行讀取每一個類別所對應的像素值['b']]
像素類別圖.append(類別所對應的RGB像素值)
像素類別圖轉為np = np.array(像素類別圖).astype('uint8')
輸出影像保存路徑 = 'E:/Camvid-FCN/test_pred_img/'
for 索引, 影像標簽資料字典 in enumerate(測驗資料):
測驗影像資料 = 影像標簽資料字典['影像'].to(計算單元)
測驗影像標簽 = 影像標簽資料字典['標簽'].long().to(計算單元)
預測 = 網路(測驗影像資料)
預測 = F.log_softmax(預測, dim=1)
預測標簽 = 預測.max(1)[1].squeeze().cpu().data.numpy()
預測標簽圖 = 像素類別圖轉為np[預測標簽]
預測標簽圖 = 預測標簽圖.squeeze()
預測標簽圖轉換格式并保存 = Image.fromarray(預測標簽圖)
預測標簽圖轉換格式并保存.save(輸出影像保存路徑 + str(索引) + '.png')
print('完成')
轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/301674.html
標籤：AI