图片分割--UNet

最新推荐文章于 2025-06-18 15:23:05 发布

原创最新推荐文章于 2025-06-18 15:23:05 发布 · 1.6k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #深度学习 #计算机视觉

1.网络结构

结构可以分为两部分左边部分是编码结构,进行特征提取右边是解码结果,进行特征还原

2.数据集准备

import os.path
from torchvision import transforms
from torch.utils.data import Dataset
from utils import *

#数据归一化
transform = transforms.Compose([
    transforms.ToTensor()
])


class MyDataset(Dataset):
    def __init__(self,path):
        self.path = path
        #获取索引的名字 E:\Pcproject\LB-UNet-main\isic2018\train
        self.name = os.listdir(os.path.join(path,'masks'))

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

    def __getitem__(self,index):
        segment_name = self.name[index]
        segment_path = os.path.join(self.path,'masks',segment_name)
        #原图地址
        image_path = os.path.join(self.path,'images',segment_name)
        #规范图片的大小尺寸
        segment_image = keep_image_size_open(segment_path)
        image = keep_image_size_open(image_path)
        return transform(image),transform(segment_image)

if __name__=='__main__':
    data = MyDataset('E:/Pcproject/pythonProjectlw/UNet')
    print(data[0][0].shape)
    print(data[0][0].shape)

数据图片规范函数:

from PIL import Image


def keep_image_size_open(path,size = (256,256)):
    #打开图像文件
    img = Image.open(path)
    #取最长边 获取图像尺寸 最长边
    temp = max(img.size)
    #创建空白图像
    mask = Image.new('RGB',(temp,temp),(0,0,0))
    #粘贴原始图像
    mask.paste(img,(0,0))
    #调整图像大小
    mask = mask.resize(size)
    #返回调整后的图像
    return mask

下面是数据的文件位置

数据集是皮肤病理分析的图片

根据unet网络结构可知有三个结构一部分是conv卷积一部分是上采样一部分是下采样

3.定义板块

1.卷积模块

from torch import nn

#卷积板块
class Conv_Block(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(Conv_Block,self).__init__()
        self.layer = nn.Sequential(
            #第一个卷积
            #padding_mode='reflect':填充的是镜像数值 比如第一行第一个数是1 第二行第1列是2 那么在第一行上面填充的值就是以1为中心对称的数字2
            #可以将填充的数值也作为特征 加强特征提取的能力
            nn.Conv2d(in_channel,out_channel,3,1,1,padding_mode='reflect',bias=False),
            nn.BatchNorm2d(out_channel),
            nn.Dropout(0.3),
            nn.LeakyReLU(),
            #第二个卷积
            nn.Conv2d(out_channel, out_channel, 3, 1, 1, padding_mode='reflect', bias=False),
            nn.BatchNorm2d(out_channel),
            nn.Dropout(0.3),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.layer(x)

padding_mode='reflect':填充的是镜像数值比如第一行第一个数是1 第二行第1列是2 那么在第一行上面填充的值就是以1为中心对称的数字2,可以将填充的数值也作为特征加强特征提取的能力

2.下采样模块

图中的的max pool最大池化进行下采样但是最大池化没有特征提取丢特征丢的太多了

#下采样模块
class DownSample(nn.Module):
    def __init__(self,channel):
        super(DownSample,self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(channel,channel,3,2,1,padding_mode='reflect',bias=False),
            nn.BatchNorm2d(channel),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.layer(x)

如果使用最大池化的化:



class DownSample(nn.Module):
    def __init__(self, channel):
        super(DownSample, self).__init__()
        self.layer = nn.Sequential(
            nn.MaxPool2d(2),  # 使用最大池化替换卷积操作
            nn.BatchNorm2d(channel),
            nn.LeakyReLU()
        )
    
    def forward(self, x):
        return self.layer(x)

3.上采样模块

使用插值法:

from torch.nn import functional as F
class UpSample(nn.Module):
    def __init__(self,channel):
        super(UpSample,self).__init__()
        self.layer = nn.Conv2d(channel,channel//2,1,1)
    def forward(self,x,feature_map):
        up = F.interpolate(x,scale_factor=2,mode='nearest')
        out = self.layer(up)
        return torch.cat((out,feature_map),dim=1)

super(UpSample, self).__init__() :
创建一个卷积层self.layer，它将输入通道数从channel减少到channel//2（即通道数减半）。
使用1x1的卷积核和步长为1，这意味着卷积层不会改变输入特征图的空间维度（高度和宽度）。
但是，它会改变通道数，因为输出通道数被设置为channel//2。
self.layer = nn.Conv2d(channel, channel//2, 1, 1)

使用F.interpolate函数对输入特征图x进行上采样。
scale_factor=2表示将特征图的高度和宽度都放大两倍。
mode='nearest'表示使用最近邻插值方法。
up = F.interpolate(x, scale_factor=2, mode='nearest')
将上采样后的特征图通过之前定义的卷积层self.layer
这将减少通道数（从原始通道数减半），同时保持（或可能稍微改变，取决于卷积层的权重初始化）空间维度。

根据结构:上采样时需要对上一个特征进行拼接:

使用torch.cat在通道维度（dim=1）上将输出特征图out和额外的特征图feature_map进行拼接。
这意味着输出特征图的通道数将是out的通道数加上feature_map的通道数。

4.定义Unet网络

class UNet(nn.Module):
    def __init__(self):
        super(UNet,self).__init__()
        #下卷积采样
        self.c1 = Conv_Block(3,64)
        self.d1 = DownSample(64)
        self.c2 = Conv_Block(64,128)
        self.d2 = DownSample(128)
        self.c3 = Conv_Block(128,256)
        self.d3 = DownSample(256)
        self.c4 = Conv_Block(256,512)
        self.d4 = DownSample(512)
        self.c5 = Conv_Block(512,1024)
        #上采样
        self.u1 = UpSample(1024)
        self.c6 = Conv_Block(1024,512)
        self.u2 = UpSample(512)
        self.c7 = Conv_Block(512,256)
        self.u3 = UpSample(256)
        self.c8 = Conv_Block(256,128)
        self.u4 = UpSample(128)
        self.c9 = Conv_Block(128,64)
        #输出
        self.out = nn.Conv2d(64,3,3,1,1)
        self.Th = nn.Sigmoid()


    def forward(self,x):
        R1 = self.c1(x)
        R2 = self.c2(self.d1(R1))
        R3 = self.c3(self.d2(R2))
        R4 = self.c4(self.d3(R3))
        R5 = self.c5(self.d4(R4))
        #拼接
        o1 = self.c6(self.u1(R5,R4))
        o2 = self.c7(self.u2(o1, R3))
        o3 = self.c8(self.u3(o2, R2))
        o4 = self.c9(self.u4(o3, R1))

        return self.Th(self.out(o4))


if __name__ == '__main__':
    x = torch.randn(2,3,256,256)
    net = UNet()
    print(net(x).shape)

5.训练代码

import os.path
from torchvision import transforms
from torch.utils.data import Dataset
from utils import *

#数据归一化
transform = transforms.Compose([
    transforms.ToTensor()
])


class MyDataset(Dataset):
    def __init__(self,path):
        self.path = path
        #获取索引的名字 E:\Pcproject\LB-UNet-main\isic2018\train
        self.name = os.listdir(os.path.join(path,'masks'))

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

    def __getitem__(self,index):
        segment_name = self.name[index]
        segment_path = os.path.join(self.path,'masks',segment_name)
        #原图地址
        image_path = os.path.join(self.path,'images',segment_name)
        #规范图片的大小尺寸
        segment_image = keep_image_size_open(segment_path)
        image = keep_image_size_open(image_path)
        return transform(image),transform(segment_image)

if __name__=='__main__':
    data = MyDataset('E:/Pcproject/pythonProjectlw/UNet')
    print(data[0][0].shape)
    print(data[0][0].shape)