【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

时间:2022-08-19 06:32:39

U-Net: Convolutional Networks for Biomedical Image Segmentation

论文作者:Olaf Ronneberger, Philipp Fischer, and Thomas Brox

源码地址:http://lmb.informatik.uni-freiburg.de/ (基于Caffe)

阅读整理:

竞赛

ISBI challenge for segmentation of neuronal structures in electron microscopic stacks

数据

30张512*512图片,主要依靠data augmentation(细胞形变等)

代码:给出的源码中可以直接运行测试单张图片,使用目前的GPU大概一秒钟。不知为何换成自己的图片去测试可能就不太好用了,可能是因为细胞过于密集?也可能是loss写的不对?论文中提到了端到端训练,代码中好像并不能直接运行实现,连基本的solver都找不到(呜呜……)。自己太菜,使用源码不如网上大神自己搭建的keras和pytorch的效果好,先总结一下论文,至于实验待日后填坑吧……不过代码中有现成的matlab借口来测试,运行起来很方便。

数据处理:

1、官方的数据集只有一张tif,最开始看到的时候整个人蒙圈了。后来才知道要自己拆分,代码参考:

from libtiff import *
imgdir = TIFF3D.open("train-labels.tif")
imgarr = imgdir.read_image()
for i in range(imgarr.shape[0]):
    savepath = "./train-labels/"+ str(i) + '.tif'
    img = TIFF.open(savepath, 'w')
img.write_image(imgarr[i])

一张tif会被拆成30张tif。

2、数据增强,对于细胞分割主要采用细胞形变,主要参考http://download.csdn.net/detail/u012931582/9817058

(C++工程,注意代码中opencv要使用2.4.10版本,版本3有一些头文件会找不到)

       3、如果使用caffe,需要生成hdf5文件。

注意:u-net很神奇,输入输出图片大小竟然不一样,记得调整label尺寸。还有,label直接读进去也不行,难免有误差,读进去的不一定是0和1的标签,搞不好就是0和0.9最后全变成0了,白训练……以下只是个生成hdf5文件的例子,####代码需要自己根据需要修改,但是仅供参考,因为我还没训练成功,姑且当个生成.h5文件的代码备份。

import random
from PIL import Image
import numpy as np
import os
import h5py
from PIL import Image

LIST_FILE = ['list_train.txt', 'list_test.txt']######################
HDF5_LIST = 'HDF5/list_hdf5.txt'##############

print '\nplease wait...'

#write
Phase='TRAIN'
slice_num=100 #every 100 img form a hdf5 file

if Phase=='TRAIN':
    image_dir=LIST_FILE[0]
elif Phase=='TEST':
    image_dir=LIST_FILE[1]
else:
    print 'error!'
    
files=[]
with open(image_dir) as f0:
    for line in f0.readlines():
        files.append(line[:-1])  
        
random.shuffle(files)#随机打乱文件顺序#############

sum_file=len(files) # sum of img
count_end=int(sum_file/slice_num) #num of hdf5 files:count_end+1


for count in range (count_end+1):
    files_part=[]
    if count==count_end:
        files_part=files[count_end*slice_num:]
    else:
        files_part=files[count*slice_num:(count+1)*slice_num]
    
    # data :eg 1channel 96*32
    datas = np.zeros((len(files_part),1, 512, 512)) ###输入tif尺寸
    # label eg 1*2
    labels = np.zeros((len(files_part),1,340, 340)) ####label尺寸

    for ii, _file in enumerate(files_part):
        train_img='./train/train_img_512/'+_file #######################
        train_label='./train/train_label_crop340/'+_file ################
        #print _file
        datas[ii, :, :] = np.array(Image.open(train_img)).astype(np.float32) / 256
        labels[ii, :, :] = \
            ((np.array(Image.open(train_label)).astype(np.float32) / 256)>0.5)\
                .astype(np.float32)####大于0.5输出1,否则输出0
                
    #New=Image.fromarray(labels[0][0])
    #New.show()

    hdf5filename=Phase+'_hdf5_'+str(count)+'.h5'
    with h5py.File('HDF5/'+hdf5filename, 'w') as f:
        f['data'] = datas
        f['label'] = labels
        f.close()
    #生成hdf5文件列表用于prototxt中
    with open(HDF5_LIST, 'a') as f:
        f.write('caffe-unet-src/cell_data/HDF5/'+hdf5filename + '\n')
        f.close()
        
    print 'hdf5 file : %d'%(count+1)
  
print '\ndone!'

网络结构

 【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

 

输出比输入图片要小。全卷积网络。像素级二分类。

预测的时候,为了检测边缘需要对输入图片进行处理:将边缘镜像。

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

对于同类密集紧邻的细胞图片:增加有细胞接触的背景部分的loss权重。

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

注意:输出的图片比输入的小,文中提到一定要选好输入尺寸,从而使得每个pooling层都能整除,才能准确分割。

 

训练:

Caffe : SGD;batch size=1;high momentum (0.99);

Energy function:

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

提前计算好训练集中每个ground truth segmentation 的weight map,计算公式如下:

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

【论文阅读】U-Net: Convolutional Networks for Biomedical Image Segmentation

参数初始化很重要:标准正态分布,方差2/N。N denotes the number of incoming nodes of one neuron . E.g. for a 3x3 convolution and 64 feature channels in the previous layer N = 9 · 64 = 576.

数据增强:平移、旋转、形变、灰度变化。形变最重要。原文叙述为We generate smooth deformations using random displacement vectors on a coarse 3 by 3 grid. The displacements are sampled from a Gaussian distribution with 10 pixels standard deviation. Per-pixel displacements are then computed using bicubic interpolation. Drop-out layers at the end of the contracting path perform further implicit data augmentation.