U-Net分割

最新推荐文章于 2025-06-27 14:57:01 发布

原创

最新推荐文章于 2025-06-27 14:57:01 发布 · 3.8k 阅读

54 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #U-Net

基于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