196 lines
6.6 KiB
Python
196 lines
6.6 KiB
Python
|
import torch
|
|||
|
|
import argparse
|
|||
|
|
import cv2
|
|||
|
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
from torch.utils.data import DataLoader
|
|||
|
|
from torch import nn, optim
|
|||
|
|
from torchvision.transforms import transforms
|
|||
|
|
from unet import Unet
|
|||
|
|
from dataset import LiverDataset
|
|||
|
|
from common_tools import transform_invert
|
|||
|
|
|
|||
|
|
|
|||
|
|
def makedir(dir):
|
|||
|
|
if not os.path.exists(dir):
|
|||
|
|
os.mkdir(dir)
|
|||
|
|
|
|||
|
|
|
|||
|
|
val_interval = 1
|
|||
|
|
# 是否使用cuda
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
|
|||
|
|
# 均为灰度图像，只需要转换为tensor
|
|||
|
|
x_transforms = transforms.ToTensor()
|
|||
|
|
y_transforms = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
train_curve = list()
|
|||
|
|
valid_curve = list()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def train_model(model, criterion, optimizer, dataload, num_epochs=100):
|
|||
|
|
makedir('./model')
|
|||
|
|
model_path = "./model/weights_100.pth"
|
|||
|
|
if os.path.exists(model_path):
|
|||
|
|
model.load_state_dict(torch.load(model_path, map_location=device))
|
|||
|
|
start_epoch = 20
|
|||
|
|
print('加载成功！')
|
|||
|
|
else:
|
|||
|
|
start_epoch = 0
|
|||
|
|
print('无保存模型，将从头开始训练！')
|
|||
|
|
|
|||
|
|
for epoch in range(start_epoch+1, num_epochs):
|
|||
|
|
print('Epoch {}/{}'.format(epoch, num_epochs))
|
|||
|
|
print('-' * 10)
|
|||
|
|
dt_size = len(dataload.dataset)
|
|||
|
|
epoch_loss = 0
|
|||
|
|
step = 0
|
|||
|
|
for x, y in dataload:
|
|||
|
|
step += 1
|
|||
|
|
inputs = x.to(device)
|
|||
|
|
labels = y.to(device)
|
|||
|
|
# zero the parameter gradients
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
# forward
|
|||
|
|
outputs = model(inputs)
|
|||
|
|
loss = criterion(outputs, labels)
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
epoch_loss += loss.item()
|
|||
|
|
train_curve.append(loss.item())
|
|||
|
|
print("%d/%d,train_loss:%0.3f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item()))
|
|||
|
|
print("epoch %d loss:%0.3f" % (epoch, epoch_loss/step))
|
|||
|
|
if (epoch + 1) % 50 == 0:
|
|||
|
|
torch.save(model.state_dict(), './model/weights_%d.pth' % (epoch + 1))
|
|||
|
|
|
|||
|
|
# Validate the model
|
|||
|
|
valid_dataset = LiverDataset("data/val", transform=x_transforms, target_transform=y_transforms)
|
|||
|
|
valid_loader = DataLoader(valid_dataset, batch_size=4, shuffle=True)
|
|||
|
|
if (epoch + 2) % val_interval == 0:
|
|||
|
|
loss_val = 0.
|
|||
|
|
model.eval()
|
|||
|
|
with torch.no_grad():
|
|||
|
|
step_val = 0
|
|||
|
|
for x, y in valid_loader:
|
|||
|
|
step_val += 1
|
|||
|
|
x = x.type(torch.FloatTensor)
|
|||
|
|
inputs = x.to(device)
|
|||
|
|
labels = y.to(device)
|
|||
|
|
outputs = model(inputs)
|
|||
|
|
loss = criterion(outputs, labels)
|
|||
|
|
loss_val += loss.item()
|
|||
|
|
|
|||
|
|
valid_curve.append(loss_val)
|
|||
|
|
print("epoch %d valid_loss:%0.3f" % (epoch, loss_val / step_val))
|
|||
|
|
|
|||
|
|
train_x = range(len(train_curve))
|
|||
|
|
train_y = train_curve
|
|||
|
|
|
|||
|
|
train_iters = len(dataload)
|
|||
|
|
valid_x = np.arange(1, len(
|
|||
|
|
valid_curve) + 1) * train_iters * val_interval # 由于valid中记录的是EpochLoss，需要对记录点进行转换到iterations
|
|||
|
|
valid_y = valid_curve
|
|||
|
|
|
|||
|
|
plt.plot(train_x, train_y, label='Train')
|
|||
|
|
plt.plot(valid_x, valid_y, label='Valid')
|
|||
|
|
|
|||
|
|
plt.legend(loc='upper right')
|
|||
|
|
plt.ylabel('loss value')
|
|||
|
|
plt.xlabel('Iteration')
|
|||
|
|
plt.show()
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 训练模型
|
|||
|
|
def train(args):
|
|||
|
|
model = Unet(1, 1).to(device)
|
|||
|
|
batch_size = args.batch_size
|
|||
|
|
criterion = nn.BCEWithLogitsLoss()
|
|||
|
|
optimizer = optim.Adam(model.parameters())
|
|||
|
|
liver_dataset = LiverDataset("./data/train", transform=x_transforms, target_transform=y_transforms)
|
|||
|
|
dataloaders = DataLoader(liver_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
|||
|
|
train_model(model, criterion, optimizer, dataloaders)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 显示模型的输出结果
|
|||
|
|
def test(args):
|
|||
|
|
model = Unet(1, 1)
|
|||
|
|
model.load_state_dict(torch.load(args.ckpt, map_location='cuda'))
|
|||
|
|
liver_dataset = LiverDataset("data/val", transform=x_transforms, target_transform=y_transforms)
|
|||
|
|
dataloaders = DataLoader(liver_dataset, batch_size=1)
|
|||
|
|
|
|||
|
|
save_root = './data/predict'
|
|||
|
|
|
|||
|
|
model.eval()
|
|||
|
|
plt.ion()
|
|||
|
|
index = 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for x, ground in dataloaders:
|
|||
|
|
x = x.type(torch.FloatTensor)
|
|||
|
|
y = model(x)
|
|||
|
|
x = torch.squeeze(x)
|
|||
|
|
x = x.unsqueeze(0)
|
|||
|
|
ground = torch.squeeze(ground)
|
|||
|
|
ground = ground.unsqueeze(0)
|
|||
|
|
img_ground = transform_invert(ground, y_transforms)
|
|||
|
|
img_x = transform_invert(x, x_transforms)
|
|||
|
|
img_y = torch.squeeze(y).numpy()
|
|||
|
|
# cv2.imshow('img', img_y)
|
|||
|
|
src_path = os.path.join(save_root, "predict_%d_s.png" % index)
|
|||
|
|
save_path = os.path.join(save_root, "predict_%d_o.png" % index)
|
|||
|
|
ground_path = os.path.join(save_root, "predict_%d_g.png" % index)
|
|||
|
|
img_ground.save(ground_path)
|
|||
|
|
# img_x.save(src_path)
|
|||
|
|
cv2.imwrite(save_path, img_y * 255)
|
|||
|
|
index = index + 1
|
|||
|
|
# plt.imshow(img_y)
|
|||
|
|
# plt.pause(0.5)
|
|||
|
|
# plt.show()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 计算Dice系数
|
|||
|
|
def dice_calc(args):
|
|||
|
|
root = './data/predict'
|
|||
|
|
nums = len(os.listdir(root)) // 3
|
|||
|
|
dice = list()
|
|||
|
|
dice_mean = 0
|
|||
|
|
for i in range(nums):
|
|||
|
|
ground_path = os.path.join(root, "predict_%d_g.png" % i)
|
|||
|
|
predict_path = os.path.join(root, "predict_%d_o.png" % i)
|
|||
|
|
img_ground = cv2.imread(ground_path)
|
|||
|
|
img_predict = cv2.imread(predict_path)
|
|||
|
|
intersec = 0
|
|||
|
|
x = 0
|
|||
|
|
y = 0
|
|||
|
|
for w in range(256):
|
|||
|
|
for h in range(256):
|
|||
|
|
intersec += img_ground.item(w, h, 1) * img_predict.item(w, h, 1) / (255 * 255)
|
|||
|
|
x += img_ground.item(w, h, 1) / 255
|
|||
|
|
y += img_predict.item(w, h, 1) / 255
|
|||
|
|
if x + y == 0:
|
|||
|
|
current_dice = 1
|
|||
|
|
else:
|
|||
|
|
current_dice = round(2 * intersec / (x + y), 3)
|
|||
|
|
dice_mean += current_dice
|
|||
|
|
dice.append(current_dice)
|
|||
|
|
dice_mean /= len(dice)
|
|||
|
|
print(dice)
|
|||
|
|
print(round(dice_mean, 3))
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
#参数解析
|
|||
|
|
parse = argparse.ArgumentParser()
|
|||
|
|
parse.add_argument("--action", type=str, help="train, test or dice", default="test")
|
|||
|
|
parse.add_argument("--batch_size", type=int, default=4)
|
|||
|
|
parse.add_argument("--ckpt", type=str, help="the path of model weight file", default="./model/weights_100.pth")
|
|||
|
|
# parse.add_argument("--ckpt", type=str, help="the path of model weight file")
|
|||
|
|
args = parse.parse_args()
|
|||
|
|
|
|||
|
|
if args.action == "train":
|
|||
|
|
train(args)
|
|||
|
|
elif args.action == "test":
|
|||
|
|
test(args)
|
|||
|
|
elif args.action == "dice":
|
|||
|
|
dice_calc(args)
|