64 lines
1.9 KiB
Python
64 lines
1.9 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
|
|||
|
|
import PIL.Image as Image
|
|||
|
|
|
|||
|
|
val_interval = 1
|
|||
|
|
# 是否使用cuda
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
|
|||
|
|
# 均为灰度图像，只需要转换为tensor
|
|||
|
|
x_transforms = transforms.ToTensor()
|
|||
|
|
y_transforms = transforms.ToTensor()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 显示模型的输出结果
|
|||
|
|
def test(args):
|
|||
|
|
model = Unet(1, 1)
|
|||
|
|
model.load_state_dict(torch.load(args.ckpt, map_location='cuda'))
|
|||
|
|
image_path = "/home/shared/wy/flask_web/Unet_liver_seg-master/data/val_data/Data/P1_T1_00038.png"
|
|||
|
|
img_x = Image.open(image_path).convert('L')
|
|||
|
|
img_x1 = x_transforms(img_x)
|
|||
|
|
dataloaders = DataLoader(img_x, batch_size=1)
|
|||
|
|
|
|||
|
|
save_root = './data/predict'
|
|||
|
|
|
|||
|
|
model.eval()
|
|||
|
|
plt.ion()
|
|||
|
|
index = 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
|
|||
|
|
x = img_x1.type(torch.FloatTensor)
|
|||
|
|
y = model(x)
|
|||
|
|
x = torch.squeeze(x)
|
|||
|
|
x = x.unsqueeze(0)
|
|||
|
|
img_x = transform_invert(x, x_transforms)
|
|||
|
|
img_y = torch.squeeze(y).numpy()
|
|||
|
|
save_path = os.path.join(save_root, "predict_%d_re.png" % index)
|
|||
|
|
cv2.imwrite(save_path, img_y * 255)
|
|||
|
|
index = index + 1
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
#参数解析
|
|||
|
|
parse = argparse.ArgumentParser()
|
|||
|
|
parse.add_argument("--action", type=str, help="train, test or dice", default="test")
|
|||
|
|
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)
|