53 lines
1.6 KiB
Python
53 lines
1.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
|
|||
|
|
import PIL.Image as Image
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
class ImageSegmentation:
|
|||
|
|
|
|||
|
|
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():
|
|||
|
|
model = Unet(1, 1)
|
|||
|
|
model.load_state_dict(torch.load('./model/weights_100.pth', map_location='cuda'))
|
|||
|
|
x_path = "/home/shared/wy/flask_web/Unet_liver_seg-master/data/val/Data/P8_T1_00070.png"
|
|||
|
|
img_x = Image.open(x_path).convert('L')
|
|||
|
|
img_x = x_transforms(img_x)
|
|||
|
|
|
|||
|
|
save_root = './data/predict'
|
|||
|
|
|
|||
|
|
# 获取当前时间
|
|||
|
|
current_time = datetime.now()
|
|||
|
|
|
|||
|
|
# 将当前时间格式化为字符串
|
|||
|
|
time_str = current_time.strftime("%Y%m%d_%H%M%S")
|
|||
|
|
model.eval()
|
|||
|
|
plt.ion()
|
|||
|
|
with torch.no_grad():
|
|||
|
|
img_x = img_x.unsqueeze(0)
|
|||
|
|
x = img_x.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_%s_.png" % time_str)
|
|||
|
|
cv2.imwrite(save_path, img_y * 255)
|