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!'
网络结构:
输出比输入图片要小。全卷积网络。像素级二分类。
预测的时候,为了检测边缘需要对输入图片进行处理:将边缘镜像。
对于同类密集紧邻的细胞图片:增加有细胞接触的背景部分的loss权重。
注意:输出的图片比输入的小,文中提到一定要选好输入尺寸,从而使得每个pooling层都能整除,才能准确分割。
训练:
Caffe : SGD;batch size=1;high momentum (0.99);
Energy function:
提前计算好训练集中每个ground truth segmentation 的weight map,计算公式如下:
参数初始化很重要:标准正态分布,方差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.