Pytorch实现UNet例子学习

参考：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)