精华内容
下载资源
问答
  • pytorch实现unet

    2021-02-06 11:11:44
    实现起来不是很复杂,代码如下: # -*- coding: utf-8 -*- import torch import torch.nn as nn class unet(nn.Module): def __init__(self): super().__init__() #conv1 self.conv1=nn.Sequential( nn.Conv...

    unet是非常经典的图像分割的网络,因网络结构形似字母U而著称

    实现起来不是很复杂,代码如下:

    # -*- coding: utf-8 -*-
    import torch
    import torch.nn as nn
    
    class unet(nn.Module):
        def __init__(self):
            super().__init__()
            #conv1
            self.conv1=nn.Sequential(
            nn.Conv2d(1,64,3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,64,3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True))
            #conv2
            self.conv2=nn.Sequential(
            nn.Conv2d(64,128,3),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128,128,3),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True))
            #conv3
            self.conv3=nn.Sequential(
            nn.Conv2d(128,256,3),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256,256,3),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True))
            #conv4
            self.conv4=nn.Sequential(
            nn.Conv2d(256,512,3),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512,512,3),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True))
            #conv5
            self.conv5=nn.Sequential(
            nn.Conv2d(1024,512,3),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512,256,3),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True))
            #conv6
            self.conv6=nn.Sequential(
            nn.Conv2d(512,256,3),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256,128,3),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True))
            #conv7
            self.conv7=nn.Sequential(
            nn.Conv2d(256,128,3),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128,64,3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True))
            
            self.trans=nn.Sequential(
            nn.Conv2d(512,1024,kernel_size=3),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024,512,3),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True))
            
            self.end=nn.Sequential(
            nn.Conv2d(128,64,kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,64,3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,2,3),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True)
            )
            
            self.unSample=nn.Upsample(mode='bilinear',scale_factor=2)
            self.pool=nn.MaxPool2d(kernel_size=2,stride=2)
            
        def forward(self,x):
            out_conv1=self.conv1(x)
            out=self.pool(out_conv1)
            out_conv2=self.conv2(out)
            out=self.pool(out_conv2)
            out_conv3=self.conv3(out)
            out=self.pool(out_conv3)
            out_conv4=self.conv4(out)
            out=self.pool(out_conv4)
            out=self.trans(out)
            out=self.unSample(out)
            out=torch.cat((out,out_conv4[:,:,:out.shape[2],:out.shape[3]]),1)
            out=self.conv5(out)
            out=self.unSample(out)
            out=torch.cat((out,out_conv3[:,:,:out.shape[2],:out.shape[3]]),1)
            out=self.conv6(out)
            out=self.unSample(out)
            out=torch.cat((out,out_conv2[:,:,:out.shape[2],:out.shape[3]]),1)        
            out=self.conv7(out)
            out=self.unSample(out)
            out=torch.cat((out,out_conv1[:,:,:out.shape[2],:out.shape[3]]),1)
            return self.end(out)
        
    if __name__=='__main__':
        input = torch.randn(1,1,572,572)
        net=unet()
        output = net(input)
        print(output)
        torch.save(net,'unet.pth')    
        torch.onnx.export(net, input,  "unet.onnx", export_params=True, opset_version=10,          
                      do_constant_folding=True,   input_names = ['input'],   output_names = ['output'], 
    )
        
    
    

    运行结果如下:

     此代码没有进行训练,是直接使用初始化权重的结果

    展开全文
  • 参考:https://github.com/milesial/Pytorch-UNet 实现的是二值汽车图像语义分割,包括 dense CRF 后处理. 使用python3,我的环境是python3.6 1.使用 1> 预测 1)查看所有的可用选项: python ...

    参考:https://github.com/milesial/Pytorch-UNet

    实现的是二值汽车图像语义分割,包括 dense CRF 后处理.

    使用python3,我的环境是python3.6

     

    1.使用

    1> 预测

    1)查看所有的可用选项:

    python predict.py -h

    返回:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ python predict.py -h
    usage: predict.py [-h] [--model FILE] --input INPUT [INPUT ...]
                      [--output INPUT [INPUT ...]] [--cpu] [--viz] [--no-save]
                      [--no-crf] [--mask-threshold MASK_THRESHOLD] [--scale SCALE]
    
    optional arguments:
      -h, --help            show this help message and exit
      --model FILE, -m FILE
                            Specify the file in which is stored the model (default
                            : 'MODEL.pth')  #指明使用的训练好的模型文件,默认使用MODEL.pth
      --input INPUT [INPUT ...], -i INPUT [INPUT ...] #指明要进行预测的图像文件,必须要有的值
                            filenames of input images
      --output INPUT [INPUT ...], -o INPUT [INPUT ...] #指明预测后生成的图像文件的名字
                            filenames of ouput images
      --cpu, -c             Do not use the cuda version of the net #指明使用CPU,默认为false,即默认使用GPU
      --viz, -v             Visualize the images as they are processed #当图像被处理时,将其可视化,默认为false,即不可以可视化
      --no-save, -n         Do not save the output masks #不存储得到的预测图像到某图像文件中,和--viz结合使用,即可对预测结果可视化,但是不存储结果,默认为false,即会保存结果
      --no-crf, -r          Do not use dense CRF postprocessing #指明不使用CRF对输出进行后处理,默认为false,即使用CRF
      --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD
                            Minimum probability value to consider a mask pixel #最小化考虑掩模像素为白色的概率值,默认为0.5
                            white
      --scale SCALE, -s SCALE
                            Scale factor for the input images #输入图像的比例因子,默认为0.5

     

    2)预测单一图片image.jpg并存储结果到output.jpg的命令

    python predict.py -i image.jpg -o output.jpg

    测试一下:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ python predict.py --cpu --viz -i image.jpg -o output.jpg
    Loading model MODEL.pth
    Using CPU version of the net, this may be very slow
    Model loaded !
    
    Predicting image image.jpg ...
    /anaconda3/envs/deeplearning/lib/python3.6/site-packages/torch/nn/modules/upsampling.py:129: UserWarning: nn.Upsample is deprecated. Use nn.functional.interpolate instead.
      warnings.warn("nn.{} is deprecated. Use nn.functional.interpolate instead.".format(self.name))
    /anaconda3/envs/deeplearning/lib/python3.6/site-packages/torch/nn/functional.py:1332: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.
      warnings.warn("nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.")
    Visualizing results for image image.jpg, close to continue ...

    返回可视化图片为:

    关闭该可视化图片命令就会运行结束:

    Mask saved to output.jpg
    (deeplearning) userdeMBP:Pytorch-UNet-master user$ 

    并且在当前文件夹中生成名为output.jpg的文件,该图为:

     

     

    3)预测多张图片并显示,预测结果不存储:

    python predict.py -i image1.jpg image2.jpg --viz --no-save

    测试:

    先得到的是image1.jpg的可视化结果:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ python predict.py -i image1.jpg image2.jpg --viz --no-save --cpu
    Loading model MODEL.pth
    Using CPU version of the net, this may be very slow
    Model loaded !
    
    Predicting image image1.jpg ...
    /anaconda3/envs/deeplearning/lib/python3.6/site-packages/torch/nn/modules/upsampling.py:129: UserWarning: nn.Upsample is deprecated. Use nn.functional.interpolate instead.
      warnings.warn("nn.{} is deprecated. Use nn.functional.interpolate instead.".format(self.name))
    /anaconda3/envs/deeplearning/lib/python3.6/site-packages/torch/nn/functional.py:1332: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.
      warnings.warn("nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.")
    Visualizing results for image image1.jpg, close to continue ...

    图为:

    关闭这个后就会接着生成image2.jpg的可视化结果:

    Predicting image image2.jpg ...
    Visualizing results for image image2.jpg, close to continue ...

    返回图为:

    这时候关闭该可视化服务就会结束了,并且没有在本地保存生成的图片

     

    4)如果你的计算机只有CPU,即CPU-only版本,使用选项--cpu指定

    5)你可以指定你使用的训练好的模型文件,使用--mode MODEL.pth

     

    6)如果使用上面的命令选项--no-crf:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ python predict.py -i image1.jpg image2.jpg --viz --no-save --cpu --no-crf

    返回的结果是:

    还有:

    可见crf后处理后,可以将一些不符合事实的判断结果给剔除,使得结果更加精确

     

    2〉训练

    python train.py -h

    首先需要安装模块pydensecrf,实现CRF条件随机场的模块:

    pip install pydensecrf
    但是出错:
    pydensecrf/densecrf/include/Eigen/Core:22:10: fatal error: 'complex' file not found #include
    <complex> ^~~~~~~~~ 1 warning and 1 error generated. error: command 'gcc' failed with exit status 1 ---------------------------------------- Failed building wheel for pydensecrf Running setup.py clean for pydensecrf Failed to build pydensecrf

    解决办法,参考https://github.com/lucasb-eyer/pydensecrf:

    先安装cython,需要0.22以上的版本:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ pip install -U cython
    Installing collected packages: cython
    Successfully installed cython-0.29.7

    然后从git安装最新版本:

    pip install git+https://github.com/lucasb-eyer/pydensecrf.git

    但还是没有成功

     

    后面找到了新的方法,使用conda来安装就成功了:

    userdeMacBook-Pro:~ user$ conda install -n deeplearning -c conda-forge pydensecrf

    -c指明从conda-forge下载模块

    conda-forge是可以安装软件包的附加渠道,使用该conda-forge频道取代defaults

    因为直接安装conda install -n deeplearning pydensecrf找不到该模块

     

    这时候运行python train.py -h可见支持的选项的信息:

    (deeplearning) userdeMBP:Pytorch-UNet-master user$ python train.py -h
    Usage: train.py [options]
    
    Options:
      -h, --help            show this help message and exit
      -e EPOCHS, --epochs=EPOCHS
                            number of epochs #指明迭代的次数
      -b BATCHSIZE, --batch-size=BATCHSIZE
                            batch size #图像批处理的大小
      -l LR, --learning-rate=LR
                            learning rate #使用的学习率
      -g, --gpu             use cuda #使用GPU进行训练
      -c LOAD, --load=LOAD  load file model #下载预训练的文件,在该基础上进行训练
      -s SCALE, --scale=SCALE
                            downscaling factor of the images #图像的缩小因子

     

    3>代码分析

    1》unet定义网络

    unet/unet_parts.py

    # sub-parts of the U-Net model
    
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    
    #实现左边的横向卷积
    class double_conv(nn.Module): 
        '''(conv => BN => ReLU) * 2'''
        def __init__(self, in_ch, out_ch):
            super(double_conv, self).__init__()
            self.conv = nn.Sequential(
                #以第一层为例进行讲解
                #输入通道数in_ch,输出通道数out_ch,卷积核设为kernal_size 3*3,padding为1,stride为1,dilation=1
                #所以图中H*W能从572*572 变为 570*570,计算为570 = ((572 + 2*padding - dilation*(kernal_size-1) -1) / stride ) +1
                nn.Conv2d(in_ch, out_ch, 3, padding=1), 
                nn.BatchNorm2d(out_ch), #进行批标准化,在训练时,该层计算每次输入的均值与方差,并进行移动平均
                nn.ReLU(inplace=True), #激活函数
                nn.Conv2d(out_ch, out_ch, 3, padding=1), #再进行一次卷积,从570*570变为 568*568
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True)
            )
    
        def forward(self, x):
            x = self.conv(x)
            return x
    
    #实现左边第一行的卷积
    class inconv(nn.Module):# 
        def __init__(self, in_ch, out_ch):
            super(inconv, self).__init__()
            self.conv = double_conv(in_ch, out_ch) # 输入通道数in_ch为3, 输出通道数out_ch为64
    
        def forward(self, x):
            x = self.conv(x)
            return x
    
    #实现左边的向下池化操作,并完成另一层的卷积
    class down(nn.Module):
        def __init__(self, in_ch, out_ch):
            super(down, self).__init__()
            self.mpconv = nn.Sequential(
                nn.MaxPool2d(2),
                double_conv(in_ch, out_ch)
            )
    
        def forward(self, x):
            x = self.mpconv(x)
            return x
    
    #实现右边的向上的采样操作,并完成该层相应的卷积操作
    class up(nn.Module): 
        def __init__(self, in_ch, out_ch, bilinear=True):
            super(up, self).__init__()
    
            #  would be a nice idea if the upsampling could be learned too,
            #  but my machine do not have enough memory to handle all those weights
            if bilinear:#声明使用的上采样方法为bilinear——双线性插值,默认使用这个值,计算方法为 floor(H*scale_factor),所以由28*28变为56*56
                self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            else: #否则就使用转置卷积来实现上采样,计算式子为 (Height-1)*stride - 2*padding -kernal_size +output_padding
                self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)
    
            self.conv = double_conv(in_ch, out_ch)
    
        def forward(self, x1, x2): #x2是左边特征提取传来的值
            #第一次上采样返回56*56,但是还没结束
            x1 = self.up(x1)
            
            # input is CHW, [0]是batch_size, [1]是通道数,更改了下,与源码不同
            diffY = x1.size()[2] - x2.size()[2] #得到图像x2与x1的H的差值,56-64=-8
            diffX = x1.size()[3] - x2.size()[3] #得到图像x2与x1的W差值,56-64=-8
    
            #用第一次上采样为例,即当上采样后的结果大小与右边的特征的结果大小不同时,通过填充来使x2的大小与x1相同
            #对图像进行填充(-4,-4,-4,-4),左右上下都缩小4,所以最后使得64*64变为56*56
            x2 = F.pad(x2, (diffX // 2, diffX - diffX//2,
                            diffY // 2, diffY - diffY//2))
            
            # for padding issues, see 
            # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
            # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
            
            #将最后上采样得到的值x1和左边特征提取的值进行拼接,dim=1即在通道数上进行拼接,由512变为1024
            x = torch.cat([x2, x1], dim=1)
            x = self.conv(x)
            return x
    
    #实现右边的最高层的最右边的卷积
    class outconv(nn.Module):
        def __init__(self, in_ch, out_ch):
            super(outconv, self).__init__()
            self.conv = nn.Conv2d(in_ch, out_ch, 1)
    
        def forward(self, x):
            x = self.conv(x)
            return x

     

    unet/unetmodel.py

    # full assembly of the sub-parts to form the complete net
    
    import torch.nn.functional as F
    
    from .unet_parts import *
    
    class UNet(nn.Module):
        def __init__(self, n_channels, n_classes): #图片的通道数,1为灰度图像,3为彩色图像
            super(UNet, self).__init__()
            self.inc = inconv(n_channels, 64) #假设输入通道数n_channels为3,输出通道数为64
            self.down1 = down(64, 128)
            self.down2 = down(128, 256)
            self.down3 = down(256, 512)
            self.down4 = down(512, 512)
            self.up1 = up(1024, 256)
            self.up2 = up(512, 128)
            self.up3 = up(256, 64)
            self.up4 = up(128, 64)
            self.outc = outconv(64, n_classes)
    
        def forward(self, x):
            x1 = self.inc(x)
            x2 = self.down1(x1)
            x3 = self.down2(x2)
            x4 = self.down3(x3)
            x5 = self.down4(x4)
            x = self.up1(x5, x4)
            x = self.up2(x, x3)
            x = self.up3(x, x2)
            x = self.up4(x, x1)
            x = self.outc(x)
            return F.sigmoid(x) #进行二分类

     

     

    2》utils

    实现dense CRF的代码utils/crf.py:

    详细可见pydensecrf的使用

    #coding:utf-8
    import numpy as np
    import pydensecrf.densecrf as dcrf
    
    def dense_crf(img, output_probs): #img为输入的图像,output_probs是经过网络预测后得到的结果
        h = output_probs.shape[0] #高度
        w = output_probs.shape[1] #宽度
    
        output_probs = np.expand_dims(output_probs, 0)
        output_probs = np.append(1 - output_probs, output_probs, axis=0)
    
        d = dcrf.DenseCRF2D(w, h, 2) #NLABELS=2两类标注,车和不是车
        U = -np.log(output_probs) #得到一元势
        U = U.reshape((2, -1)) #NLABELS=2两类标注
        U = np.ascontiguousarray(U) #返回一个地址连续的数组
        img = np.ascontiguousarray(img)
    
        d.setUnaryEnergy(U) #设置一元势
    
        d.addPairwiseGaussian(sxy=20, compat=3) #设置二元势中高斯情况的值
        d.addPairwiseBilateral(sxy=30, srgb=20, rgbim=img, compat=10)#设置二元势众双边情况的值
    
        Q = d.inference(5) #迭代5次推理
        Q = np.argmax(np.array(Q), axis=0).reshape((h, w)) #得列中最大值的索引结果
    
        return Q

     

    utils/utils.py

    import random
    import numpy as np
    
    #将图像分成左右两块
    def get_square(img, pos): 
        """Extract a left or a right square from ndarray shape : (H, W, C))"""
        h = img.shape[0]
        if pos == 0:
            return img[:, :h]
        else:
            return img[:, -h:]
    
    def split_img_into_squares(img):
        return get_square(img, 0), get_square(img, 1)
    
    #对图像进行转置,将(H, W, C)变为(C, H, W)
    def hwc_to_chw(img):
        return np.transpose(img, axes=[2, 0, 1])
    
    def resize_and_crop(pilimg, scale=0.5, final_height=None):
        w = pilimg.size[0] #得到图片的宽
        h = pilimg.size[1]#得到图片的高
        #默认scale为0.5,即将高和宽都缩小一半
        newW = int(w * scale) 
        newH = int(h * scale)
    
        #如果没有指明希望得到的最终高度
        if not final_height:
            diff = 0
        else:
            diff = newH - final_height
        #重新设定图片的大小
        img = pilimg.resize((newW, newH))
        #crop((left,upper,right,lower))函数,从图像中提取出某个矩形大小的图像。它接收一个四元素的元组作为参数,各元素为(left, upper, right, lower),坐标系统的原点(0, 0)是左上角
        #如果没有设置final_height,其实就是取整个图片
        #如果设置了final_height,就是取一个上下切掉diff // 2,最后高度为final_height的图片
        img = img.crop((0, diff // 2, newW, newH - diff // 2))
        return np.array(img, dtype=np.float32)
    
    def batch(iterable, batch_size):
        """批量处理列表"""
        b = []
        for i, t in enumerate(iterable):
            b.append(t)
            if (i + 1) % batch_size == 0:
                yield b
                b = []
    
        if len(b) > 0:
            yield b
    
    #然后将数据分为训练集和验证集两份
    def split_train_val(dataset, val_percent=0.05):
        dataset = list(dataset)
        length = len(dataset) #得到数据集大小
        n = int(length * val_percent) #验证集的数量
        random.shuffle(dataset) #将数据打乱
        return {'train': dataset[:-n], 'val': dataset[-n:]} 
    
    #对像素值进行归一化,由[0,255]变为[0,1]
    def normalize(x):
        return x / 255
    
    #将两个图片合并起来
    def merge_masks(img1, img2, full_w):
        h = img1.shape[0]
    
        new = np.zeros((h, full_w), np.float32)
        new[:, :full_w // 2 + 1] = img1[:, :full_w // 2 + 1]
        new[:, full_w // 2 + 1:] = img2[:, -(full_w // 2 - 1):]
    
        return new
    
    
    # credits to https://stackoverflow.com/users/6076729/manuel-lagunas
    def rle_encode(mask_image):
        pixels = mask_image.flatten()
        # We avoid issues with '1' at the start or end (at the corners of
        # the original image) by setting those pixels to '0' explicitly.
        # We do not expect these to be non-zero for an accurate mask,
        # so this should not harm the score.
        pixels[0] = 0
        pixels[-1] = 0
        runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
        runs[1::2] = runs[1::2] - runs[:-1:2]
        return runs

     

    utils/data_vis.py实现结果的可视化:

    import matplotlib.pyplot as plt
    
    def plot_img_and_mask(img, mask):
        fig = plt.figure()
        a = fig.add_subplot(1, 2, 1) #先是打印输入的图片
        a.set_title('Input image')
        plt.imshow(img)
    
        b = fig.add_subplot(1, 2, 2) #然后打印预测得到的结果图片
        b.set_title('Output mask')
        plt.imshow(mask)
        plt.show()

     

    utils/load.py

    #
    # load.py : utils on generators / lists of ids to transform from strings to
    #           cropped images and masks
    
    import os
    
    import numpy as np
    from PIL import Image
    
    from .utils import resize_and_crop, get_square, normalize, hwc_to_chw
    
    
    def get_ids(dir):
        """返回目录中的id列表"""
        return (f[:-4] for f in os.listdir(dir)) #图片名字的后4位为数字,能作为图片id
    
    
    def split_ids(ids, n=2):
        """将每个id拆分为n个,为每个id创建n个元组(id, k)"""
        #等价于for id in ids:
        #       for i in range(n):
        #           (id, i)
        #得到元祖列表[(id1,0),(id1,1),(id2,0),(id2,1),...,(idn,0),(idn,1)]
        #这样的作用是后面会通过后面的0,1作为utils.py中get_square函数的pos参数,pos=0的取左边的部分,pos=1的取右边的部分
        return ((id, i)  for id in ids for i in range(n))
    
    
    def to_cropped_imgs(ids, dir, suffix, scale):
        """从元组列表中返回经过剪裁的正确img"""
        for id, pos in ids:
            im = resize_and_crop(Image.open(dir + id + suffix), scale=scale) #重新设置图片大小为原来的scale倍
            yield get_square(im, pos) #然后根据pos选择图片的左边或右边
    
    def get_imgs_and_masks(ids, dir_img, dir_mask, scale):
        """返回所有组(img, mask)"""
    
        imgs = to_cropped_imgs(ids, dir_img, '.jpg', scale)
    
        # need to transform from HWC to CHW
        imgs_switched = map(hwc_to_chw, imgs) #对图像进行转置,将(H, W, C)变为(C, H, W)
        imgs_normalized = map(normalize, imgs_switched) #对像素值进行归一化,由[0,255]变为[0,1]
    
        masks = to_cropped_imgs(ids, dir_mask, '_mask.gif', scale) #对图像的结果也进行相同的处理
    
        return zip(imgs_normalized, masks) #并将两个结果打包在一起
    
    
    def get_full_img_and_mask(id, dir_img, dir_mask):
        im = Image.open(dir_img + id + '.jpg')
        mask = Image.open(dir_mask + id + '_mask.gif')
        return np.array(im), np.array(mask)

     

    3》预测

    predict.py使用训练好的U-net网络对图像进行预测,使用dense CRF进行后处理:

    #coding:utf-8
    import argparse
    import os
    
    import numpy as np
    import torch
    import torch.nn.functional as F
    
    from PIL import Image
    
    from unet import UNet
    from utils import resize_and_crop, normalize, split_img_into_squares, hwc_to_chw, merge_masks, dense_crf
    from utils import plot_img_and_mask
    
    from torchvision import transforms
    
    def predict_img(net,
                    full_img,
                    scale_factor=0.5,
                    out_threshold=0.5,
                    use_dense_crf=True,
                    use_gpu=False):
    
        net.eval() #进入网络的验证模式,这时网络已经训练好了
        img_height = full_img.size[1] #得到图片的高
        img_width = full_img.size[0] #得到图片的宽
    
        img = resize_and_crop(full_img, scale=scale_factor) #在utils文件夹的utils.py中定义的函数,重新定义图像大小并进行切割,然后将图像转为数组np.array
        img = normalize(img) #对像素值进行归一化,由[0,255]变为[0,1]
    
        left_square, right_square = split_img_into_squares(img)#将图像分成左右两块,来分别进行判断
    
        left_square = hwc_to_chw(left_square) #对图像进行转置,将(H, W, C)变为(C, H, W),便于后面计算
        right_square = hwc_to_chw(right_square)
    
        X_left = torch.from_numpy(left_square).unsqueeze(0) #将(C, H, W)变为(1, C, H, W),因为网络中的输入格式第一个还有一个batch_size的值
        X_right = torch.from_numpy(right_square).unsqueeze(0)
        
        if use_gpu:
            X_left = X_left.cuda()
            X_right = X_right.cuda()
    
        with torch.no_grad(): #不计算梯度
            output_left = net(X_left)
            output_right = net(X_right)
    
            left_probs = output_left.squeeze(0)
            right_probs = output_right.squeeze(0)
    
            tf = transforms.Compose(
                [
                    transforms.ToPILImage(), #重新变成图片
                    transforms.Resize(img_height), #恢复原来的大小
                    transforms.ToTensor() #然后再变成Tensor格式
                ]
            )
            
            left_probs = tf(left_probs.cpu())
            right_probs = tf(right_probs.cpu())
    
            left_mask_np = left_probs.squeeze().cpu().numpy()
            right_mask_np = right_probs.squeeze().cpu().numpy()
    
        full_mask = merge_masks(left_mask_np, right_mask_np, img_width)#将左右两个拆分后的图片合并起来
    
        #对得到的结果根据设置决定是否进行CRF处理
        if use_dense_crf:
            full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
    
        return full_mask > out_threshold
    
    
    
    def get_args():
        parser = argparse.ArgumentParser()
        parser.add_argument('--model', '-m', default='MODEL.pth', #指明使用的训练好的模型文件,默认使用MODEL.pth
                            metavar='FILE',
                            help="Specify the file in which is stored the model"
                                 " (default : 'MODEL.pth')")
        parser.add_argument('--input', '-i', metavar='INPUT', nargs='+',  #指明要进行预测的图像文件
                            help='filenames of input images', required=True)
    
        parser.add_argument('--output', '-o', metavar='INPUT', nargs='+', #指明预测后生成的图像文件的名字
                            help='filenames of ouput images')
        parser.add_argument('--cpu', '-c', action='store_true', #指明使用CPU
                            help="Do not use the cuda version of the net",
                            default=False)
        parser.add_argument('--viz', '-v', action='store_true', 
                            help="Visualize the images as they are processed", #当图像被处理时,将其可视化
                            default=False)
        parser.add_argument('--no-save', '-n', action='store_true', #不存储得到的预测图像到某图像文件中,和--viz结合使用,即可对预测结果可视化,但是不存储结果
                            help="Do not save the output masks",
                            default=False)
        parser.add_argument('--no-crf', '-r', action='store_true', #指明不使用CRF对输出进行后处理
                            help="Do not use dense CRF postprocessing",
                            default=False)
        parser.add_argument('--mask-threshold', '-t', type=float, 
                            help="Minimum probability value to consider a mask pixel white", #最小概率值考虑掩模像素为白色
                            default=0.5)
        parser.add_argument('--scale', '-s', type=float,
                            help="Scale factor for the input images", #输入图像的比例因子
                            default=0.5)
    
        return parser.parse_args()
    
    def get_output_filenames(args):#从输入的选项args值中得到输出文件名
        in_files = args.input 
        out_files = []
    
        if not args.output: #如果在选项中没有指定输出的图片文件的名字,那么就会根据输入图片文件名,在其后面添加'_OUT'后缀来作为输出图片文件名
            for f in in_files:
                pathsplit = os.path.splitext(f) #将文件名和扩展名分开,pathsplit[0]是文件名,pathsplit[1]是扩展名
                out_files.append("{}_OUT{}".format(pathsplit[0], pathsplit[1])) #得到输出图片文件名
        elif len(in_files) != len(args.output): #如果设置了output名,查看input和output的数量是否相同,即如果input是两张图,那么设置的output也必须是两个,否则报错
            print("Error : Input files and output files are not of the same length")
            raise SystemExit()
        else:
            out_files = args.output
    
        return out_files
    
    def mask_to_image(mask):
        return Image.fromarray((mask * 255).astype(np.uint8)) #从数组array转成Image
    
    if __name__ == "__main__":
        args = get_args() #得到输入的选项设置的值
        in_files = args.input #得到输入的图像文件
        out_files = get_output_filenames(args) #从输入的选项args值中得到输出文件名
    
        net = UNet(n_channels=3, n_classes=1) #定义使用的model为UNet,调用在UNet文件夹下定义的unet_model.py,定义图像的通道为3,即彩色图像,判断类型设为1种
    
        print("Loading model {}".format(args.model)) #指定使用的训练好的model
    
        if not args.cpu: #指明使用GPU
            print("Using CUDA version of the net, prepare your GPU !")
            net.cuda()
            net.load_state_dict(torch.load(args.model))
        else: #否则使用CPU
            net.cpu()
            net.load_state_dict(torch.load(args.model, map_location='cpu'))
            print("Using CPU version of the net, this may be very slow")
    
        print("Model loaded !")
    
        for i, fn in enumerate(in_files): #对图片进行预测
            print("\nPredicting image {} ...".format(fn))
    
            img = Image.open(fn)
            if img.size[0] < img.size[1]: #(W, H, C)
                print("Error: image height larger than the width")
    
            mask = predict_img(net=net,
                               full_img=img,
                               scale_factor=args.scale,
                               out_threshold=args.mask_threshold,
                               use_dense_crf= not args.no_crf,
                               use_gpu=not args.cpu)
    
            if args.viz: #可视化输入的图片和生成的预测图片
                print("Visualizing results for image {}, close to continue ...".format(fn))
                plot_img_and_mask(img, mask)
    
            if not args.no_save:#设置为False,则保存
                out_fn = out_files[i]
                result = mask_to_image(mask) #从数组array转成Image
                result.save(out_files[i]) #然后保存
    
                print("Mask saved to {}".format(out_files[i]))

     

    4》训练

    import sys
    import os
    from optparse import OptionParser
    import numpy as np
    
    import torch
    import torch.backends.cudnn as cudnn
    import torch.nn as nn
    from torch import optim
    
    from eval import eval_net
    from unet import UNet
    from utils import get_ids, split_ids, split_train_val, get_imgs_and_masks, batch
    
    def train_net(net,
                  epochs=5,
                  batch_size=1,
                  lr=0.1,
                  val_percent=0.05,
                  save_cp=True,
                  gpu=False,
                  img_scale=0.5):
    
        dir_img = 'data/train/' #训练图像文件夹
        dir_mask = 'data/train_masks/' #图像的结果文件夹
        dir_checkpoint = 'checkpoints/' #训练好的网络保存文件夹
    
        ids = get_ids(dir_img)#图片名字的后4位为数字,能作为图片id
    
        #得到元祖列表为[(id1,0),(id1,1),(id2,0),(id2,1),...,(idn,0),(idn,1)]
        #这样的作用是后面重新设置生成器时会通过后面的0,1作为utils.py中get_square函数的pos参数,pos=0的取左边的部分,pos=1的取右边的部分
        #这样图片的数量就会变成2倍
        ids = split_ids(ids) 
    
        iddataset = split_train_val(ids, val_percent) #将数据分为训练集和验证集两份
    
        print('''
        Starting training:
            Epochs: {}
            Batch size: {}
            Learning rate: {}
            Training size: {}
            Validation size: {}
            Checkpoints: {}
            CUDA: {}
        '''.format(epochs, batch_size, lr, len(iddataset['train']),
                   len(iddataset['val']), str(save_cp), str(gpu)))
    
        N_train = len(iddataset['train']) #训练集长度
    
        optimizer = optim.SGD(net.parameters(), #定义优化器
                              lr=lr,
                              momentum=0.9,
                              weight_decay=0.0005)
    
        criterion = nn.BCELoss()#损失函数
    
        for epoch in range(epochs): #开始训练
            print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
            net.train() #设置为训练模式
    
            # reset the generators重新设置生成器
            # 对输入图片dir_img和结果图片dir_mask进行相同的图片处理,即缩小、裁剪、转置、归一化后,将两个结合在一起,返回(imgs_normalized, masks)
            train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale)
            val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale)
    
            epoch_loss = 0
    
            for i, b in enumerate(batch(train, batch_size)):
                imgs = np.array([i[0] for i in b]).astype(np.float32) #得到输入图像数据
                true_masks = np.array([i[1] for i in b]) #得到图像结果数据
    
                imgs = torch.from_numpy(imgs)
                true_masks = torch.from_numpy(true_masks)
    
                if gpu:
                    imgs = imgs.cuda()
                    true_masks = true_masks.cuda()
    
                masks_pred = net(imgs) #图像输入的网络后得到结果masks_pred,结果为灰度图像
                masks_probs_flat = masks_pred.view(-1) #将结果压扁
    
                true_masks_flat = true_masks.view(-1) 
    
                loss = criterion(masks_probs_flat, true_masks_flat) #对两个结果计算损失
                epoch_loss += loss.item()
    
                print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
    
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
    
            print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) #一次迭代后得到的平均损失
    
            if 1:
                val_dice = eval_net(net, val, gpu)
                print('Validation Dice Coeff: {}'.format(val_dice))
    
            if save_cp:
                torch.save(net.state_dict(),
                           dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
                print('Checkpoint {} saved !'.format(epoch + 1))
    
    
    
    def get_args():
        parser = OptionParser()
        parser.add_option('-e', '--epochs', dest='epochs', default=5, type='int', #设置迭代数
                          help='number of epochs')
        parser.add_option('-b', '--batch-size', dest='batchsize', default=10, #设置训练批处理数
                          type='int', help='batch size')
        parser.add_option('-l', '--learning-rate', dest='lr', default=0.1, #设置学习率
                          type='float', help='learning rate')
        parser.add_option('-g', '--gpu', action='store_true', dest='gpu', #是否使用GPU,默认是不使用
                          default=False, help='use cuda')
        parser.add_option('-c', '--load', dest='load', #下载之前预训练好的模型
                          default=False, help='load file model')
        parser.add_option('-s', '--scale', dest='scale', type='float', #图像的缩小因子,用来重新设置图片大小
                          default=0.5, help='downscaling factor of the images') 
    
        (options, args) = parser.parse_args()
        return options
    
    if __name__ == '__main__':
        args = get_args() #得到设置的所有参数信息
    
        net = UNet(n_channels=3, n_classes=1)
    
        if args.load: #是否加载预先训练好的模型
            net.load_state_dict(torch.load(args.load))
            print('Model loaded from {}'.format(args.load))
    
        if args.gpu: #是否使用GPU,设置为True,则使用
            net.cuda()
            # cudnn.benchmark = True # faster convolutions, but more memory
    
        try: #开始训练
            train_net(net=net,
                      epochs=args.epochs,
                      batch_size=args.batchsize,
                      lr=args.lr,
                      gpu=args.gpu,
                      img_scale=args.scale)
        except KeyboardInterrupt: #如果键盘输入ctrl+c停止,则会将结果保存在INTERRUPTED.pth中
            torch.save(net.state_dict(), 'INTERRUPTED.pth')
            print('Saved interrupt')
            try:
                sys.exit(0)
            except SystemExit:
                os._exit(0)

     

    转载于:https://www.cnblogs.com/wanghui-garcia/p/10719121.html

    展开全文
  • 该存储库包含U-Net和FCN的简单PyTorch实现,这是Ronneberger等人提出的深度学习细分方法。 和龙等。 用于训练的合成图像/遮罩 首先克隆存储库并cd到项目目录。 import matplotlib . pyplot as plt import numpy as...
  • pytorch实现unet网络,专门用于进行图像分割训练。该代码打过kaggle上的 Carvana Image Masking Challenge from a high definition image.
  • Unet:U-Net: Convolutional Networks for Biomedical Image Segmentation目标检测模型在Pytorch当中的实现 目录 性能情况 Performance 所需环境 Environment 注意事项 Attention 文件下载 Download 预测步骤 How...
  • 基于pytorchUNet分割网络demo实现,及训练自己的数据集。包括对相关报错的分析。收集了几个比较好的前辈的网址。
  • pytorchUnet实现医学图像分割

    万次阅读 多人点赞 2019-07-23 11:03:57
    MIT License Copyright (c) 2019 PJ-Javis 目录: 链接: https://pan.baidu.com/s/1WVVo6VDv5NkjgdYfkIO7Fw 提取... ...unet.py import torch.nn as nn import torch from torch import autograd class Double...

    MIT License

    Copyright (c) 2019 PJ-Javis

    目录:

    链接: https://pan.baidu.com/s/1WVVo6VDv5NkjgdYfkIO7Fw 提取码: 14mi

     

    unet.py

    import torch.nn as nn
    import torch
    from torch import autograd
    
    class DoubleConv(nn.Module):
        def __init__(self, in_ch, out_ch):
            super(DoubleConv, self).__init__()
            self.conv = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                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

    dataset.py

    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)
            mask=os.path.join(root,"%03d_mask.png"%i)
            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)

    main.py

    import numpy as np
    import torch
    import argparse
    from torch.utils.data import DataLoader
    from torch import autograd, optim
    from torchvision.transforms import transforms
    from unet import Unet
    from dataset import LiverDataset
    
    
    # 是否使用cuda
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 把多个步骤整合到一起, channel=(channel-mean)/std, 因为是分别对三个通道处理
    x_transforms = transforms.Compose([
        transforms.ToTensor(),  # -> [0,1]
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # ->[-1,1]
    ])
    
    # 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("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)
    
    
    # 显示模型的输出结果
    def test():
        model = Unet(3, 1)
        model.load_state_dict(torch.load(args.ckp,map_location='cpu'))
        liver_dataset = LiverDataset("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 x, _ in dataloaders:
                y=model(x)
                img_y=torch.squeeze(y).numpy()
                plt.imshow(img_y)
                plt.pause(0.01)
            plt.show()
    
    
    parse = argparse.ArgumentParser()
    # parse.add_argument("action", type=str, help="train or test")
    parse.add_argument("--batch_size", type=int, default=1)
    parse.add_argument("--ckp", type=str, help="the path of model weight file")
    args = parse.parse_args()
    
    # train
    #train()
    
    # test()
    args.ckp = "weights_19.pth"
    test()

    测试结果:

     

    展开全文
  • 使用PyTorchUNet实现人像匹配。 细分演示结果: 哑光演示结果: 为了方便演示,我通过Flask构建了API服务,最后将其部署在微信小程序中。 微信小程序的代码部分在此处。 依存关系 Python 3.6 PyTorch> = 1.1.0...
  • PyTorchUnet网络实现脑肿瘤图像分割

    千次阅读 2019-10-24 13:01:25
    本文主要利用pytorch实现了U-Net网络的脑肿瘤分割,并展示了部分分割效果图。希望对你有所帮助! 2 源代码 (1)网络结构代码 import torch.nn as nn import torch from torch import autograd class DoubleConv...

    1 介绍

    U-Net是一篇基本结构非常好的论文,主要是针对生物医学图片的分割,而且,在今后的许多对医学图像的分割网络中,很大一部分会采取U-Net作为网络的主干。相对于当年的,在EM segmentation challenge at ISBI 2012上做到比当时的best更好。而且速度也非常的快。其有一个很好的优点,就是在小数据集上也是能做得比较好的。就比如EM 2012这个数据集就只是30个果蝇第一龄幼虫腹侧神经所索的连续部分透射电子显微镜图。

    本文主要利用pytorch实现了U-Net网络的脑肿瘤分割,并展示了部分分割效果图。希望对你有所帮助!

    2 源代码

    (1)网络结构代码

    import torch.nn as nn
    import torch
    from torch import autograd
    
    class DoubleConv(nn.Module):
        def __init__(self, in_ch, out_ch):
            super(DoubleConv, self).__init__()
            self.conv = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                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
    

    (2)数据集准备

    import torch.utils.data as data
    import PIL.Image as Image
    import os
    
    
    def make_dataset(rootdata,roottarget):#获取img和mask的地址
        imgs = []
        filename_data = [x for x in os.listdir(rootdata)]
        for name in filename_data:
            img = os.path.join(rootdata, name)
            mask = os.path.join(roottarget, name)
            imgs.append((img, mask))#作为元组返回
        return imgs
    
    
    class MyDataset(data.Dataset):
        def __init__(self, rootdata, roottarget, transform=None, target_transform=None):
            imgs = make_dataset(rootdata,roottarget)
            #print(imgs)
            self.imgs = imgs
            self.transform = transform
            self.target_transform = target_transform
    
        def __getitem__(self, index):
            x_path, y_path = self.imgs[index]
            # print(x_path)
            img_x = Image.open(x_path).convert('L')#读取并转换为二值图像
            img_y = Image.open(y_path).convert('L')
            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)
            # print(img_x.shape[0])
            # print(img_x.shape[1])
            # print(img_x.shape[2])
            # print(img_x)
            return img_x, img_y
    
        def __len__(self):
            return len(self.imgs)

     (3)训练与测试

    import numpy as np
    import torch
    import argparse
    from torch.utils.data import DataLoader
    from torch import autograd, optim
    from torch.autograd import Variable
    from torchvision import transforms
    from unet import Unet
    from dataset import MyDataset
    from dataset import make_dataset
    import os
    import cv2
    import time
    
    # 是否使用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()
    
    
    def train_model(model, criterion, optimizer, dataload, num_epochs=150):
        for epoch in range(0,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.6f" %
                      (step,(dt_size - 1) // dataload.batch_size + 1, loss.item()))
            print("epoch %d loss:%0.6f" % (epoch, epoch_loss))
        torch.save(model.cpu().state_dict(), 'weights_%d.pth' % epoch)
        torch.save(model.cpu(),'weights_%d_dc.pth' % epoch)
        return model
    
    
    #训练模型
    def train(train_data_path,train_gt_path):
        batch_size = 1
        # liver_dataset = MyDataset(
        #     "image/train/data", "image/train/gt",transform=x_transforms, target_transform=y_transforms)
        liver_dataset = MyDataset(
            train_data_path, train_gt_path, transform=x_transforms, target_transform=y_transforms)
        dataloaders = DataLoader(
            liver_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
        train_model(model, criterion, optimizer, dataloaders)
    
    
    #显示模型的输出结果
    def test(test_data_path,test_gt_path,save_pre_path):
        # liver_dataset = MyDataset(
        #     "image/val/data", "image/val/gt", transform=x_transforms, target_transform=y_transforms)
        liver_dataset = MyDataset(
                 test_data_path, test_gt_path, transform=x_transforms, target_transform=y_transforms)
        imgs = make_dataset(test_data_path, test_gt_path)
        # print(imgs[0][1])
        print(liver_dataset)
        dataloaders = DataLoader(liver_dataset, batch_size=1)
        import matplotlib.pyplot as plt
        plt.ion()
        with torch.no_grad():
            count = 0
            for x, _ in dataloaders:
                start = time.clock()
                x = x.to('cuda')
                y = model(x)
                img_y = torch.squeeze(y).cpu().numpy()
                elapsed = (time.clock() - start)
                print("Time used:",elapsed)
                plt.imsave(os.path.join(save_pre_path,os.path.basename(imgs[count][1])), img_y)
                count+=1
            plt.show()
    
    def test_forDirsImages(source_data_path,source_gt_path,save_path):
        if not os.path.exists(source_data_path):
            return
        if not os.path.exists(source_gt_path):
            return
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        sour_data_path = source_data_path
        sour_gt_path = source_gt_path
        sav_path = save_path
        for i in range(1,100):
            source_data_path = os.path.join(source_data_path,str(i))
            source_gt_path = os.path.join(source_gt_path,str(i))
            save_path = os.path.join(save_path, str(i))
            if not os.path.exists(save_path):
                os.makedirs(save_path)
            test(source_data_path, source_gt_path, save_path)
    
            source_data_path = sour_data_path
            source_gt_path = sour_gt_path
            save_path = sav_path
    
    if __name__ == '__main__':
        pretrained = False
        model = Unet(1, 1).to(device)
        if pretrained:
            model.load_state_dict(torch.load('./weights_149.pth'))
        criterion = torch.nn.BCELoss()
        optimizer = optim.Adam(model.parameters())
        train("I:\\0\\train20191022\\img","I:\\0\\train20191022\\gt") 

     

    3 效果展示

    说明:左图为原始图像,右图为分割结果

    展开全文
  • 为了用pytorch实现pointnet,需要先学习pytorch的一个网络,然后再进行实现 关于pytorch的安装,个人建议初学者安装cpu的最好,因为gpu的pytorch在windows下特别难以安装成功 因为gpu目前实在版本混乱,而且跟...
  • 原始论文 U-Net: Convolutional Networks for Biomedical Image Segmentation:点击查看 网络结构 注意事项 论文中进行... 引用参考 https://zhuanlan.zhihu.com/p/87593567 https://github.com/milesial/Pytorch-UNet
  • pytorch实现简易版unet

    千次阅读 2018-07-21 13:43:20
    最近看了一些semantic segmentation的paper,参考网上的教程(https://tuatini.me/practical-image-segmentation-with-unet/),用pytorch实现一个简化的unet练练手。 整个project的文件结构如下: 关于数据集:...
  • 基于pytorch编写Unet进行细胞结构分割

    千次阅读 2020-01-06 20:53:07
    Unet想必大家都听过它的大名,在医学图像分割方向非常著名的一篇论文,我也是怀着好奇心,以其为我想对分割任务有个初步认识为动机,实现了一下论文中说的那个挑战赛(训练数据少,训练时间短)...
  • UNetPytorch实现

    千次阅读 2020-10-07 16:29:49
    UNet的pytorch实现 原文 本文实现 训练过的UNet参数文件 提取码:1zom 1.概述 UNet是医学图像分割领域经典的论文,因其结构像字母U得名。 倘若了解过Encoder-Decoder结构、实现过DenseNet,那么实现Unet并非难事。 ...
  • 主要介绍了使用pytorch实现论文中的unet网络,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
  • 损失函数matlab代码3D-UNet-PyTorch-实现 这是Özgün Çiçek等人提出的3D UNet实现,详情请参考:。 使用的数据集:,我使用的数据集已经被其他人处理过,由于某些原因我不能在这里分享它。 如果您想在工作中...
  • Unet图像分割网络Pytorch实现

    万次阅读 多人点赞 2019-02-12 02:34:04
    之前计划写一篇tensorflow实现的,但是最近一个月接触了一下Pytorch,个人认为Pytorch相较于Tensorflow来说好用很多。本文的内容是我对Unet论文的总结与提炼,需要提醒的是,Unet原文发布的时候还没有提出BN(Batch ...
  • UNet(UNet网络的三个实现:大同小异 全是pytorch实现),其中UNet文件中包含了原始的分割数据集
  • 3D-UNetPytorch实现

    千次阅读 2020-12-08 10:05:16
    3D-UNetPytorch实现 本文主要介绍3DUNet网络,及其在LiTS2017肝脏肿瘤数据集上训练的Pytorch实现代码。 GitHub地址: 添加链接描述 LiTS2017数据集 链接: 添加链接描述 提取码:hfl8 (++||…==’’。。。_) 一....
  • Unet通常应用到单类别的语义分割,经过调整后该代码适合于多类别的语义分割。对应博客:https://blog.csdn.net/brf_UCAS/article/details/112383722
  • Unetpytorch实现

    2020-11-26 20:35:47
    Unet的网络结构如图 原论文点击这里 # -*- coding: utf-8 -*- # @Time : 2020/11/26 19:05 # @Author : Refrain_ouc # @FileName: unet_parts.py # @Software: PyCharm import torch import torch.nn.functional ...
  • UNet是医学图像处理方面著名的图像分割网络,过程是这样的:输入是一幅图,输出是目标的分割结果。继续简化就是,一幅图,编码,或者说降采样,然后解码,也就是升采样,然后输出一个分割结果。根据结果和真实分割的...

空空如也

空空如也

1 2 3 4 5 ... 8
收藏数 157
精华内容 62
关键字:

pytorch实现unet