1、先建立unet
文章中所有路径需要改成本地数据集位置
unet_parts.py
import torch
import torch.nn as nn
import torch.nn.functional as F
# 定义一个双卷积块，包括两个卷积层，归一化和ReLU激活函数
def __init__(self, in_channels, out_channels):
super().__init__()
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),#归一化处理
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
# 定义一个下采样模块，包括最大池化和双卷积块
class Down(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
# 定义一个上采样模块，包括上采样操作和双卷积块
class Up(nn.Module):
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
# 如果使用双线性插值，则使用正常卷积降低通道数
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
# 定义一个输出卷积模块，包括一个卷积层
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)unet_model
import torch.nn.functional as F
from unet_parts import *
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=True):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
#原尺寸图大小[512,512,3]
self.inc = DoubleConv(n_channels, 64)#[512,512,64]
self.down1 = Down(64, 128)#maxpooling后[256,256,64]，在经过两次卷积[256,256,128]
self.down2 = Down(128, 256)#maxpooling后[128,128,128]，在经过两次卷积[128,128,256]
self.down3 = Down(256, 512)#maxpooling后[64,64,256]，在经过两次卷积[64,64,512]
self.down4 = Down(512, 512)#maxpooling后[32,32,512]，在经过两次卷积[32,32,512]
self.up1 = Up(1024, 256, bilinear)
self.up2 = Up(512, 128, bilinear)
self.up3 = Up(256, 64, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_classes)
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)#[256,256,128]
x3 = self.down2(x2)#[128,128,256]
x4 = self.down3(x3)#[64,64,512]
x5 = self.down4(x4)#[32,32,512]
x = self.up1(x5, x4)#先将[32,32,512]上采样--->[64,64,512],与x5融合[64,64,1024]
x = self.up2(x, x3)#先将[64,64,1024]
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = self.outc(x)
return logits
if __name__ == '__main__':
net = UNet(n_channels=3, n_classes=1)
print(net)2、训练数据模型（需要时间较长）
train.py
from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch
from tqdm import tqdm
def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
# 加载训练集
isbi_dataset = ISBI_Loader(data_path)
per_epoch_num = len(isbi_dataset) / batch_size
train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
batch_size=batch_size,
shuffle=False)
# 定义RMSprop算法
optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
# 定义Loss算法
criterion = nn.BCEWithLogitsLoss()
# best_loss统计，初始化为正无穷
best_loss = float('inf')
# 训练epochs次
with tqdm(total=epochs*per_epoch_num) as pbar:
for epoch in range(epochs):
# 训练模式
net.train()
# 按照batch_size开始训练
for image, label in train_loader:
optimizer.zero_grad()
# 将数据拷贝到device中
image = image.to(device=device, dtype=torch.float32)
label = label.to(device=device, dtype=torch.float32)
# 使用网络参数，输出预测结果
pred = net(image)
# 计算loss
loss = criterion(pred, label)
# print('{}/{}：Loss/train'.format(epoch + 1, epochs), loss.item())
# 保存loss值最小的网络参数
if loss < best_loss:
best_loss = loss
torch.save(net.state_dict(), 'best_model.pth')
# 更新参数
loss.backward()
optimizer.step()
pbar.update(1)
if __name__ == "__main__":
# 选择设备，有cuda用cuda，没有就用cpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载网络，图片单通道1，分类为1。
net = UNet(n_channels=1, n_classes=1) # todo edit input_channels n_classes
# 将网络拷贝到deivce中
net.to(device=device)
# 指定训练集地址，开始训练
data_path = r"路径" # todo 修改为你本地的数据集位置
print("进度条出现卡着不动不是程序问题，是他正在计算，请耐心等待")
train_net(net, device, data_path, epochs=40, batch_size=1)
3、测试模型
test.py
import os
from tqdm import tqdm
from utils.utils_metrics import compute_mIoU, show_results
import glob
import numpy as np
import torch
import os
import cv2
from model.unet_model import UNet
def cal_miou(test_dir=r"测试路径",
pred_dir=r"保存路径（保存训练结果）", gt_dir=r"标签路径"):
miou_mode = 0
num_classes = 2
name_classes = ["background", "skin"]
if miou_mode == 0 or miou_mode == 1:
if not os.path.exists(pred_dir):
os.makedirs(pred_dir)
print("Load model.")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载网络，图片单通道，分类为1。
net = UNet(n_channels=1, n_classes=1)
# 将网络拷贝到deivce中
net.to(device=device)
# 加载模型参数
net.load_state_dict(torch.load('best_model.pth', map_location=device)) # todo
# 测试模式
net.eval()
print("Load model done.")
img_names = os.listdir(test_dir)
image_ids = [image_name.split(".")[0] for image_name in img_names]
print("Get predict result.")
for image_id in tqdm(image_ids):
image_path = os.path.join(test_dir, image_id + ".jpg")
img = cv2.imread(image_path)
origin_shape = img.shape
# print(origin_shape)
# 转为灰度图
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
img = cv2.resize(img, (512, 512))
# 转为batch为1，通道为1，大小为512*512的数组
img = img.reshape(1, 1, img.shape[0], img.shape[1])
# 转为tensor
img_tensor = torch.from_numpy(img)
# 将tensor拷贝到device中，只用cpu就是拷贝到cpu中，用cuda就是拷贝到cuda中。
img_tensor = img_tensor.to(device=device, dtype=torch.float32)
# 预测
pred = net(img_tensor)
# 提取结果
pred = np.array(pred.data.cpu()[0])[0]
pred[pred >= 0.5] = 255
pred[pred < 0.5] = 0
pred = cv2.resize(pred, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_NEAREST)
cv2.imwrite(os.path.join(pred_dir, image_id + ".png"), pred)
print("Get predict result done.")
if miou_mode == 0 or miou_mode == 2:
print("Get miou.")
print(gt_dir)
print(pred_dir)
print(num_classes)
print(name_classes)
hist, IoUs, PA_Recall, Precision = compute_mIoU(gt_dir, pred_dir, image_ids, num_classes,
name_classes) # 执行计算mIoU的函数
print("Get miou done.")
miou_out_path = "results/"
show_results(miou_out_path, hist, IoUs, PA_Recall, Precision, name_classes)
if __name__ == '__main__':
cal_miou()4、环境
5、数据集可礼貌咨询
扫一扫
4993
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
