U-Net分割

最新推荐文章于 2024-05-28 21:20:00 发布

朴素.无恙

最新推荐文章于 2024-05-28 21:20:00 发布

阅读量3.6k

点赞数 7

分类专栏： pytorch 文章标签： pytorch U-Net

本文链接：https://blog.csdn.net/weixin_40123108/article/details/88314427

版权

pytorch 专栏收录该内容

36 篇文章 10 订阅

订阅专栏

基于pytorch的U-Net分割的实现

导入一些需要的库

import numpy as np
import torch
import argparse
from torch.utils.data import DataLoader
from torch import autograd, optim
from torchvision.transforms import transforms
import torch.nn as nn
import scipy.misc
import torch.utils.data as data
import PIL.Image as Image
import os

1. 制作数据集
该数据集中data由原图像与标记（mask）好的图像组成，该数据集的形式如下所示，下述数据集的准备即将原始数据与mask数据用列表储存，即x,y=[original,mask]形式
在这里插入图片描述

#读取数据的路径
def make_dataset(root):
    imgs=[]
    n=len(os.listdir(root))//2#即共有多少张原始图片
    for i in range(n):
        img=os.path.join(root,"%03d.png"%i)#找到00i.png的路径
        mask=os.path.join(root,"%03d_mask.png"%i)#找到00i_mask.png的路径
        imgs.append((img,mask))#添加至列表
    return imgs


class LiverDataset(data.Dataset):
    def __init__(self, root, transform=None, target_transform=None):
        imgs = make_dataset(root)
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        x_path, y_path = self.imgs[index]
        img_x = Image.open(x_path)
        img_y = Image.open(y_path)
        if self.transform is not None:
            img_x = self.transform(img_x)
        if self.target_transform is not None:
            img_y = self.target_transform(img_y)
        return img_x, img_y

    def __len__(self):
        return len(self.imgs)

在这里插入图片描述

2. 定义U–Net网络模型

class DoubleConv(nn.Module):#为U_Net模型中的双卷积结构
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),#此处包含padding，为了使输出图像与原图像大小相同
            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)
        out = nn.Sigmoid()(c10)
        return out

# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

3. 模型的预处理

x_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# mask只需要转换为tensor
y_transforms = transforms.ToTensor()

#参数解析
parse=argparse.ArgumentParser()

4. 训练模型函数

def train_model(model, criterion, optimizer, dataload, num_epochs=20):
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        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()
            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))
    torch.save(model.state_dict(), 'weights_%d.pth' % epoch)
    return model

#训练模型
def train():
    model = Unet(3, 1).to(device)
    batch_size = args.batch_size
    criterion = torch.nn.BCELoss()
    optimizer = optim.Adam(model.parameters())
    liver_dataset = LiverDataset("/home/dell/Desktop/dj/U-Net/pytorch_u_net/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)
save_dir="/home/dell/Desktop/"#预测图像的存储路径

5. 模型的测试函数

#显示模型的输出结果
def test():
    model = Unet(3, 1)
    model.load_state_dict(torch.load('weights_19.pth'))
    liver_dataset = LiverDataset("/home/dell/Desktop/dj/U-Net/pytorch_u_net/data/val", transform=x_transforms,target_transform=y_transforms)
    dataloaders = DataLoader(liver_dataset, batch_size=1)
    model.eval()
    import matplotlib.pyplot as plt
    plt.ion()
    with torch.no_grad():
        for i, data in enumerate(dataloaders):
            x,z=data
            y=model(x)
            img_y=torch.squeeze(y).numpy()
            plt.subplot(1,2,1)
            z=torch.squeeze(z).numpy()
            plt.imshow(z),plt.axis("on") 
            plt.subplot(1,2,2)
            plt.imshow(img_y),plt.axis("on") 
            plt.pause(0.01)
            filename=save_dir + 'new_%d.jpg' % i#保存文件的格式 
            #print(filename) 
            scipy.misc.toimage(img_y,cmin=0.0,cmax=1.0).save(filename)#保存图像

6. 主程序

if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument("action", type=str, help="train or test")
    parse.add_argument("--batch_size", type=int, default=4)#此处根据自己电脑的性能选取batch
    
    args = parse.parse_args()

    if args.action=="train":
        train()
    elif args.action=="test":
        test()

模型的结果
在这里插入图片描述
完整程序(U_Net.py)如下：

import numpy as np
import torch
import argparse
from torch.utils.data import DataLoader
from torch import autograd, optim
from torchvision.transforms import transform
import torch.nn as nn
import scipy.misc
import torch.utils.data as data
import PIL.Image as Image
import os


#读取数据的路径
def make_dataset(root):
    imgs=[]
    n=len(os.listdir(root))//2#即共有多少张原始图片
    for i in range(n):
        img=os.path.join(root,"%03d.png"%i)#找到00i.png的路径
        mask=os.path.join(root,"%03d_mask.png"%i)#找到00i_mask.png的路径
        imgs.append((img,mask))#添加至列表
    return imgs


class LiverDataset(data.Dataset):
    def __init__(self, root, transform=None, target_transform=None):
        imgs = make_dataset(root)
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        x_path, y_path = self.imgs[index]
        img_x = Image.open(x_path)
        img_y = Image.open(y_path)
        if self.transform is not None:
            img_x = self.transform(img_x)
        if self.target_transform is not None:
            img_y = self.target_transform(img_y)
        return img_x, img_y

    def __len__(self):
        return len(self.imgs)

class DoubleConv(nn.Module):#为U_Net模型中的双卷积结构
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),#此处包含padding，为了使输出图像与原图像大小相同
            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)
        out = nn.Sigmoid()(c10)
        return out

# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# mask只需要转换为tensor
y_transforms = transforms.ToTensor()

#参数解析
parse=argparse.ArgumentParser()

def train_model(model, criterion, optimizer, dataload, num_epochs=20):
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        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()
            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))
    torch.save(model.state_dict(), 'weights_%d.pth' % epoch)
    return model

#训练模型
def train():
    model = Unet(3, 1).to(device)
    batch_size = args.batch_size
    criterion = torch.nn.BCELoss()
    optimizer = optim.Adam(model.parameters())
    liver_dataset = LiverDataset("/home/dell/Desktop/dj/U-Net/pytorch_u_net/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)

save_dir="/home/dell/Desktop/"#预测图像的存储路径


#显示模型的输出结果
def test():
    model = Unet(3, 1)
    model.load_state_dict(torch.load('weights_19.pth'))
    liver_dataset = LiverDataset("/home/dell/Desktop/dj/U-Net/pytorch_u_net/data/val", transform=x_transforms,target_transform=y_transforms)
    dataloaders = DataLoader(liver_dataset, batch_size=1)
    model.eval()
    import matplotlib.pyplot as plt
    plt.ion()
    with torch.no_grad():
        for i, data in enumerate(dataloaders):
            x,z=data
            y=model(x)
            img_y=torch.squeeze(y).numpy()
            plt.subplot(1,2,1)
            z=torch.squeeze(z).numpy()
            plt.imshow(z),plt.axis("on") 
            plt.subplot(1,2,2)
            plt.imshow(img_y),plt.axis("on") 
            plt.pause(0.01)
            filename=save_dir + 'new_%d.jpg' % i#保存文件的格式 
            #print(filename) 
            scipy.misc.toimage(img_y,cmin=0.0,cmax=1.0).save(filename)#保存图像 

if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument("action", type=str, help="train or test")
    parse.add_argument("--batch_size", type=int, default=4)#此处根据自己电脑的性能选取batch
    
    args = parse.parse_args()

    if args.action=="train":
        train()
    elif args.action=="test":
        test()

如何调用该程序，
训练阶段：进入U_Net.py的路径下执行python U_Net.py train
测试阶段：进入U_Net.py的路径下执行python U_Net.py test

朴素.无恙

关注

7
点赞
踩
54

收藏

觉得还不错? 一键收藏
45
评论
U-Net分割

U-Net分割pytorch的实现import numpy as npimport torchimport argparsefrom torch.utils.data import DataLoaderfrom torch import autograd, optimfrom torchvision.transforms import transformsfrom unet impo...
复制链接

扫一扫