unet 学习笔记-3 优化代码、调节代码结构
这篇博文是unet学习笔记的第3个，在网络上下载的代码一般要么不完整，要么太完整、庞大。对于初期学习比较吃力，写此博文，作为笔记。
unet网络结构代码，这个没变，unet.py
import torch
from torch import nn
class DoubleConv(nn.Module):
def __init__(self, in_ch, out_ch):
super(DoubleConv, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True)
)
def forward(self, input):
return self.conv(input)
class Unet(nn.Module):
def __init__(self,in_ch,out_ch):
super(Unet, self).__init__()
self.conv1 = DoubleConv(in_ch, 64)
self.pool1 = nn.MaxPool2d(2)
self.conv2 = DoubleConv(64, 128)
self.pool2 = nn.MaxPool2d(2)
self.conv3 = DoubleConv(128, 256)
self.pool3 = nn.MaxPool2d(2)
self.conv4 = DoubleConv(256, 512)
self.pool4 = nn.MaxPool2d(2)
self.conv5 = DoubleConv(512, 1024)
self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
self.conv6 = DoubleConv(1024, 512)
self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2)
self.conv7 = DoubleConv(512, 256)
self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.conv8 = DoubleConv(256, 128)
self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.conv9 = DoubleConv(128, 64)
self.conv10 = nn.Conv2d(64, out_ch, 1)
def forward(self, x):
c1 = self.conv1(x)
p1 = self.pool1(c1)
c2 = self.conv2(p1)
p2 = self.pool2(c2)
c3 = self.conv3(p2)
p3 = self.pool3(c3)
c4 = self.conv4(p3)
p4 = self.pool4(c4)
c5 = self.conv5(p4)
up_6 = self.up6(c5)
merge6 = torch.cat([up_6, c4], dim=1)
c6=self.conv6(merge6)
up_7 = self.up7(c6)
merge7 = torch.cat([up_7, c3], dim=1)
c7 = self.conv7(merge7)
up_8 = self.up8(c7)
merge8 = torch.cat([up_8, c2], dim=1)
c8 = self.conv8(merge8)
up_9 = self.up9(c8)
merge9 = torch.cat([up_9, c1], dim=1)
c9 = self.conv9(merge9)
c10 = self.conv10(c9)
return c10网络训练train.py
# -*- coding: utf-8 -*-
import torch
from torchvision.transforms import transforms
from torch import nn, optim
from unet import Unet
import numpy as np
# from tqdm import tqdm
import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
path = '/home/kitty/BiSeNet-master/data/tongue'
train_image_path = os.path.join(path, 'train')
train_label_path = os.path.join(path, 'train_labels')
test_image_path = os.path.join(path, 'test')
test_label_path = os.path.join(path, 'test_labels')
train_image = os.listdir(train_image_path)
train_label = os.listdir(train_label_path)
test_image = os.listdir(test_image_path)
test_label = os.listdir(test_label_path)
PATH = './unet_model.pt'
# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x_transforms = transforms.Compose([
transforms.ToPILImage(),
# transforms.Resize(512,512),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# mask 只需要转换为tensor
y_transforms = transforms.Compose([
transforms.ToPILImage(),
# transforms.Resize(512,512),
transforms.ToTensor(),
])
def train_model(model, criterion, optimizer, num_epochs=10):
best_model = model
min_loss = 1000
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs-1))
epoch_loss = 0
step = 0
for i in range(len(train_image)):
step += 1
image = cv2.imread(train_image_path + '/' + train_image[i], cv2.IMREAD_COLOR)
image1 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image1 = cv2.resize(image1, (512, 512))
label = cv2.imread(train_label_path + '/' + train_label[i], cv2.IMREAD_GRAYSCALE)
# label1 = cv2.cvtColor(label, cv2.COLOR_BGR2RGB)
label1 = cv2.resize(label, (512, 512))
inputs = x_transforms(image1).unsqueeze(0).to(device)
labels = y_transforms(label1).unsqueeze(0).to(device)
# inputs = x.to(device)
# labels = y.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print("%d, train_loss:%0.3f" % (step, loss.item()))
print("epoch %d loss:%0.3f" % (epoch, epoch_loss/step))
if (epoch_loss/step) < min_loss:
min_loss = (epoch_loss/step)
best_model = model
torch.save(best_model.state_dict(), PATH)
return best_model
# 训练模型
def train():
model = Unet(3,1).to(device)
# batch_size = 1
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters())
# train_dataset = TrainDataset("", "", transform=x_transforms, target_transform=y_transforms())
# dataloaders = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
train_model(model, criterion, optimizer)
def test():
model = Unet(3,1)
model.load_state_dict(torch.load(PATH))
with torch.no_grad():
for i in range(1):
image = cv2.imread(test_image_path + '/' + test_image[i], cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (512, 512))
print('image:', image.shape)
inputs = x_transforms(image).unsqueeze(0)
print('inputs:', inputs.shape)
y = model(inputs)
y = y.squeeze(0)
y = y.permute(1,2,0)
y = torch.sigmoid(y).numpy()
# print(y.min())
# print(y.max())
y = (y * 255).astype(np.uint8)
return y
if __name__ == '__main__':
print("开始训练")
# train()
print("训练完成，保存模型")
print("-"*20)
print("开始预测")
y = test()
cv2.imshow('tt', y)
cv2.waitKey(1000)
测试代码：test.py
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as transforms
from unet import Unet
path = '/home/kitty/BiSeNet-master/data/tongue'
test_image_path = os.path.join(path, 'test')
test_label_path = os.path.join(path, 'test_labels')
x_transforms = transforms.Compose([
transforms.ToPILImage(),
# transforms.Resize(512, 512),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
mopel_path = 'unet_model.pt'
def test(test_image_path, model_path):
test_image = os.listdir(test_image_path)
test_label = os.listdir(test_label_path)
# print(test_image)
# print(test_label)
# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Unet(3,1).to(device)
model.load_state_dict(torch.load(model_path))
# print(model)
with torch.no_grad():
# for i in range(len(train_image)):
for i in range(8):
image = cv2.imread(test_image_path + '/' + test_image[i], cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (512, 512))
print('image:', image.shape)
inputs = x_transforms(image).unsqueeze(0).to(device)
print('inputs.shape:', inputs.shape)
import time
print(time.time())
y = model(inputs)
print(time.time())
y = y.squeeze(0)
y = y.permute(1, 2, 0)
# print('min,max', y.min(), y.max())
y = torch.sigmoid(y)
# print('min,max', y.min(), y.max())
# print('min,max', y.min()*255, y.max()*255)
y = y.cpu().numpy()
y = (y * 255).astype(np.uint8)
# my remote ubuntu desktop, opencv-python cannot exit correctly.
# cv2.imshow('y', y)
# cv2.waitKey(3000)
# cv2.destroyAllWindows()
# plt.imshow(y, cmap=plt.get_cmap('gray'))
# y = cv2.cvtColor(y, cv2.COLOR_GRAY2RGB)
# y = plt.imshow(y)
# plt.show()
# return y
if __name__ == '__main__':
print("开始预测")
test(test_image_path, mopel_path)