Unet:实现多通道输入的医学图像分割

前言：

作为一个小白加入实验室后寒假需要自己动手搭建一个神经网络来完成脑部肿瘤的分割识别，网上几个阅读量比较大的博客都讲的比较详细，但是发现基本上都是灰度图输入，灰度图输出，而我的数据集为彩色图片要输出二值化（黑白）预测图片，而且很多的博客代码运行会有很多问题，通过学习了几个博主的博客和研究了几份代码之后完成自己的第一个完整图像分割编写，这边文章主要是针对彩图数据集输入，二分类，输出黑白图片来表示分割，这里参考的博主的博客如下：

 

主要结构：

U-net：运行你的第一个U-net进行图像分割

U-net入门--纠错过程

因为这里的网络我更改并适用于我的数据集并不能很好的运行，通过几天的调试发现问题还是在网络中，但是因为网络是黑盒运行，我也没有太好的解决方法，所以把网络结构更改为api调用，运行结构不错。

unet模型部分的api调用参考：

全卷积神经网络FCN模型Unet 图像分割项目实践

建议：

因为我的代码主要基于第一个链接的博客进行改动，所以建议先下载该博主的github代码，并阅读纠错过程先对部分进行改动，然后先尝试自行分析代码，如果是小白切记不要心浮气躁，不然代码一行都读不进去....（个人经验），如果实在读不懂，再来对比看我的分析与改动，这样至少都代码有一定熟悉度

 

数据处理：

# 创建训练数据
	def create_train_data(self):
		i = 0
		print('-'*30)
		print('Creating training images...')
		print('-'*30)
		imgs = glob.glob(self.data_path+"/*."+self.img_type)
		imgs.sort(key=lambda x: int(x.split('/')[3][:-4]))
		print(len(imgs))

		imgdatas = np.ndarray((len(imgs),self.out_rows,self.out_cols,3), dtype=np.uint8)
		imglabels = np.ndarray((len(imgs),self.out_rows,self.out_cols,1), dtype=np.uint8)
		for imgname in imgs:
			midname = imgname[imgname.rindex("/")+1:]
			img = load_img(self.data_path + "/" + midname, color_mode='rgb')
			label = load_img(self.label_path + "/" + midname, grayscale=True)
			img = img_to_array(img)
			label = img_to_array(label)
			imgdatas[i] = img
			imglabels[i] = label
			if i % 100 == 0:
				print('Done: {0}/{1} images'.format(i, len(imgs)))
			i += 1
		print('loading done')
		np.save(self.npy_path + '/imgs_train.npy', imgdatas)
		np.save(self.npy_path + '/imgs_mask_train.npy', imglabels)
		print('Saving to .npy files done.')

解读：

先是通过glob.glob读入所有训练图片的路径，这个glob的读入顺序可能并不是所想的那样（这个坑花费了我很长时间），所以我加入了一行sort函数，按照我的图片格式重新排序，如果不重新排序预测出来的mask就会对应不上，导致模型训练成功但是预测结果感觉是错的...（亲测经验）这里sort和自己的图片存储路径位置还有命名格式有关，建议自己学习相关用法后写自己的sort函数。

之后的imgdatas和imglabels通过np.ndarray开了两个空存储空间，因为彩图和灰度图的通道数不同，所以第二处改动就是imgdatas的np.ndarray函数需要更改1为3，因为这里输入的是训练彩图为三通道，之后通过for循环依次将图片读入并添加到imgdatas和imglabels这两个存储空间中。

第三处改动为img的color_mode，因为读入的是彩图当然要用彩图格式

之后的创建测试数据代码类似，同样改动三处，在此不赘述

 

# 加载训练图片与mask
	def load_train_data(self):
		print('-'*30)
		print('load train images...')
		print('-'*30)
		imgs_train = np.load(self.npy_path+"/imgs_train.npy")
		imgs_mask_train = np.load(self.npy_path+"/imgs_mask_train.npy")
		imgs_train = imgs_train.astype('float32')
		imgs_mask_train = imgs_mask_train.astype('float32')
		imgs_train /= 255
		mean = imgs_train.mean(axis = 0)
		imgs_train -= mean
		imgs_mask_train /= 255
		imgs_mask_train[imgs_mask_train > 0.5] = 1
		imgs_mask_train[imgs_mask_train <= 0.5] = 0
		return imgs_train,imgs_mask_train

# 加载测试图片
	def load_test_data(self):
		print('-'*30)
		print('load test images...')
		print('-'*30)
		imgs_test = np.load(self.npy_path+"/imgs_test.npy")
		imgs_test = imgs_test.astype('float32')
		imgs_test /= 255
		mean = imgs_test.mean(axis = 0)
		imgs_test -= mean
		return imgs_test

 

加载图片中无改动，这里主要注意下这里的二值化操作

 

if __name__ == "__main__":
# 以下注释掉的部分为数据增强代码，通过他们可以将数据进行增强

	#aug = myAugmentation()
	#aug.Augmentation()
	#aug.splitMerge()
	#aug.splitTransform()

	mydata = dataProcess(64,64)
	mydata.create_train_data()
	mydata.create_test_data()

	imgs_train,imgs_mask_train = mydata.load_train_data()
	print (imgs_train.shape,imgs_mask_train.shape)

main函数中的处理函数当然要改一下图片尺寸，这里根据自己的数据集改动即可

	def train(self):
		print("loading data")
		imgs_train, imgs_mask_train, imgs_test = self.load_data()
		print("loading data done")
		#model = self.get_unet()
		model = Unet('resnet34', input_shape=(64, 64, 3), encoder_weights=None)
		model.compile('Adam', loss='binary_crossentropy', metrics=['accuracy'])
		#model = load_model("final_model")
		print("got unet")
		model_checkpoint = ModelCheckpoint('final_model', monitor='loss',verbose=1, save_best_only=True)
		print('Fitting model...')
		model.fit(imgs_train, imgs_mask_train, batch_size=16, nb_epoch=15, verbose=1,validation_split=0.2, shuffle=True,callbacks=[model_checkpoint])
		print('predict test data')
		imgs_mask_test = model.predict(imgs_test, batch_size=1, verbose=1)
		np.save('../results/imgs_mask_test.npy', imgs_mask_test)

训练模型部分，我是直接调用的unet api，需要添加头文件，具体我会把代码放到博客的最后，input_shape改为我们的模型大小，我的是64x64的彩图，所以是64x64的大小，3通道。训练时用model.fit，可以更改batch_size和epoch来寻找最合适的参数，如果训练好的模型想多次预测不同的图片，注释掉这一行，将load_model的注释去掉直接使用训练好的模型，具体的函数参数的意义建议不会的同学一个一个google学习一下，对理解代码很有帮助，毕竟做这个不就是想在实践中学习知识么，所以不妨多花点时间学一下，个人亲测不要半小时就能都学会。

 

	def save_img(self):
		print("array to image")
		imgs = np.load('../results/imgs_mask_test.npy')
		imgs[imgs>0.5]=1
		imgs[imgs<=0.5]=0
		for i in range(imgs.shape[0]):
			img = imgs[i]
			img = array_to_img(img)
			img.save("../results/%d.tif"%(i+1))

这里我加了对结果就行二值化操作，具体看自己对结果的要求，是需要灰度图就行还是需要二值化的图片来决定需不需要二值化，主要这里不需要再乘255，这也是我踩过的一个坑，个人认为因为是float类型，所以array_to_img已经封装了将数据乘以255变成像素的操作了，最后还要要注意图片的保存格式，和训练时的一致即可

 

运行结果：

因为是api直接调的...准确度等指标都还不错

 

这里只说明了几个有改动的模块，本人水平有限也才入门学习，有问题欢迎留言交流，我会尽量回复，完整代码我放到了github，如果对你有所帮助希望star支持 ：）

GitHub代码