unet.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273
  1. # coding: utf8
  2. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. from __future__ import absolute_import
  16. from __future__ import division
  17. from __future__ import print_function
  18. from collections import OrderedDict
  19. import paddle.fluid as fluid
  20. from .model_utils.libs import scope, name_scope
  21. from .model_utils.libs import bn, bn_relu, relu
  22. from .model_utils.libs import conv, max_pool, deconv
  23. from .model_utils.libs import sigmoid_to_softmax
  24. from .model_utils.loss import softmax_with_loss
  25. from .model_utils.loss import dice_loss
  26. from .model_utils.loss import bce_loss
  27. import paddlex.utils.logging as logging
  28. class UNet(object):
  29. """实现Unet模型
  30. `"U-Net: Convolutional Networks for Biomedical Image Segmentation"
  31. <https://arxiv.org/abs/1505.04597>`
  32. Args:
  33. num_classes (int): 类别数
  34. mode (str): 网络运行模式,根据mode构建网络的输入和返回。
  35. 当mode为'train'时,输入为image(-1, 3, -1, -1)和label (-1, 1, -1, -1) 返回loss。
  36. 当mode为'train'时,输入为image (-1, 3, -1, -1)和label (-1, 1, -1, -1),返回loss,
  37. pred (与网络输入label 相同大小的预测结果,值代表相应的类别),label,mask(非忽略值的mask,
  38. 与label相同大小,bool类型)。
  39. 当mode为'test'时,输入为image(-1, 3, -1, -1)返回pred (-1, 1, -1, -1)和
  40. logit (-1, num_classes, -1, -1) 通道维上代表每一类的概率值。
  41. upsample_mode (str): UNet decode时采用的上采样方式,取值为'bilinear'时利用双线行差值进行上菜样,
  42. 当输入其他选项时则利用反卷积进行上菜样,默认为'bilinear'。
  43. use_bce_loss (bool): 是否使用bce loss作为网络的损失函数,只能用于两类分割。可与dice loss同时使用。
  44. use_dice_loss (bool): 是否使用dice loss作为网络的损失函数,只能用于两类分割,可与bce loss同时使用。
  45. 当use_bce_loss和use_dice_loss都为False时,使用交叉熵损失函数。
  46. class_weight (list/str): 交叉熵损失函数各类损失的权重。当class_weight为list的时候,长度应为
  47. num_classes。当class_weight为str时, weight.lower()应为'dynamic',这时会根据每一轮各类像素的比重
  48. 自行计算相应的权重,每一类的权重为:每类的比例 * num_classes。class_weight取默认值None是,各类的权重1,
  49. 即平时使用的交叉熵损失函数。
  50. ignore_index (int): label上忽略的值,label为ignore_index的像素不参与损失函数的计算。
  51. Raises:
  52. ValueError: use_bce_loss或use_dice_loss为真且num_calsses > 2。
  53. ValueError: class_weight为list, 但长度不等于num_class。
  54. class_weight为str, 但class_weight.low()不等于dynamic。
  55. TypeError: class_weight不为None时,其类型不是list或str。
  56. """
  57. def __init__(self,
  58. num_classes,
  59. mode='train',
  60. upsample_mode='bilinear',
  61. use_bce_loss=False,
  62. use_dice_loss=False,
  63. class_weight=None,
  64. ignore_index=255):
  65. # dice_loss或bce_loss只适用两类分割中
  66. if num_classes > 2 and (use_bce_loss or use_dice_loss):
  67. raise Exception(
  68. "dice loss and bce loss is only applicable to binary classfication"
  69. )
  70. if class_weight is not None:
  71. if isinstance(class_weight, list):
  72. if len(class_weight) != num_classes:
  73. raise ValueError(
  74. "Length of class_weight should be equal to number of classes"
  75. )
  76. elif isinstance(class_weight, str):
  77. if class_weight.lower() != 'dynamic':
  78. raise ValueError(
  79. "if class_weight is string, must be dynamic!")
  80. else:
  81. raise TypeError(
  82. 'Expect class_weight is a list or string but receive {}'.
  83. format(type(class_weight)))
  84. self.num_classes = num_classes
  85. self.mode = mode
  86. self.upsample_mode = upsample_mode
  87. self.use_bce_loss = use_bce_loss
  88. self.use_dice_loss = use_dice_loss
  89. self.class_weight = class_weight
  90. self.ignore_index = ignore_index
  91. def _double_conv(self, data, out_ch):
  92. param_attr = fluid.ParamAttr(
  93. name='weights',
  94. regularizer=fluid.regularizer.L2DecayRegularizer(
  95. regularization_coeff=0.0),
  96. initializer=fluid.initializer.TruncatedNormal(loc=0.0, scale=0.33))
  97. with scope("conv0"):
  98. data = bn_relu(
  99. conv(
  100. data,
  101. out_ch,
  102. 3,
  103. stride=1,
  104. padding=1,
  105. param_attr=param_attr))
  106. with scope("conv1"):
  107. data = bn_relu(
  108. conv(
  109. data,
  110. out_ch,
  111. 3,
  112. stride=1,
  113. padding=1,
  114. param_attr=param_attr))
  115. return data
  116. def _down(self, data, out_ch):
  117. # 下采样:max_pool + 2个卷积
  118. with scope("down"):
  119. data = max_pool(data, 2, 2, 0)
  120. data = self._double_conv(data, out_ch)
  121. return data
  122. def _up(self, data, short_cut, out_ch):
  123. # 上采样:data上采样(resize或deconv), 并与short_cut concat
  124. param_attr = fluid.ParamAttr(
  125. name='weights',
  126. regularizer=fluid.regularizer.L2DecayRegularizer(
  127. regularization_coeff=0.0),
  128. initializer=fluid.initializer.XavierInitializer(),
  129. )
  130. with scope("up"):
  131. if self.upsample_mode == 'bilinear':
  132. short_cut_shape = fluid.layers.shape(short_cut)
  133. data = fluid.layers.resize_bilinear(data, short_cut_shape[2:])
  134. else:
  135. data = deconv(
  136. data,
  137. out_ch // 2,
  138. filter_size=2,
  139. stride=2,
  140. padding=0,
  141. param_attr=param_attr)
  142. data = fluid.layers.concat([data, short_cut], axis=1)
  143. data = self._double_conv(data, out_ch)
  144. return data
  145. def _encode(self, data):
  146. # 编码器设置
  147. short_cuts = []
  148. with scope("encode"):
  149. with scope("block1"):
  150. data = self._double_conv(data, 64)
  151. short_cuts.append(data)
  152. with scope("block2"):
  153. data = self._down(data, 128)
  154. short_cuts.append(data)
  155. with scope("block3"):
  156. data = self._down(data, 256)
  157. short_cuts.append(data)
  158. with scope("block4"):
  159. data = self._down(data, 512)
  160. short_cuts.append(data)
  161. with scope("block5"):
  162. data = self._down(data, 512)
  163. return data, short_cuts
  164. def _decode(self, data, short_cuts):
  165. # 解码器设置,与编码器对称
  166. with scope("decode"):
  167. with scope("decode1"):
  168. data = self._up(data, short_cuts[3], 256)
  169. with scope("decode2"):
  170. data = self._up(data, short_cuts[2], 128)
  171. with scope("decode3"):
  172. data = self._up(data, short_cuts[1], 64)
  173. with scope("decode4"):
  174. data = self._up(data, short_cuts[0], 64)
  175. return data
  176. def _get_logit(self, data, num_classes):
  177. # 根据类别数设置最后一个卷积层输出
  178. param_attr = fluid.ParamAttr(
  179. name='weights',
  180. regularizer=fluid.regularizer.L2DecayRegularizer(
  181. regularization_coeff=0.0),
  182. initializer=fluid.initializer.TruncatedNormal(loc=0.0, scale=0.01))
  183. with scope("logit"):
  184. data = conv(
  185. data,
  186. num_classes,
  187. 3,
  188. stride=1,
  189. padding=1,
  190. param_attr=param_attr)
  191. return data
  192. def _get_loss(self, logit, label, mask):
  193. avg_loss = 0
  194. if not (self.use_dice_loss or self.use_bce_loss):
  195. avg_loss += softmax_with_loss(
  196. logit,
  197. label,
  198. mask,
  199. num_classes=self.num_classes,
  200. weight=self.class_weight,
  201. ignore_index=self.ignore_index)
  202. else:
  203. if self.use_dice_loss:
  204. avg_loss += dice_loss(logit, label, mask)
  205. if self.use_bce_loss:
  206. avg_loss += bce_loss(
  207. logit, label, mask, ignore_index=self.ignore_index)
  208. return avg_loss
  209. def generate_inputs(self):
  210. inputs = OrderedDict()
  211. inputs['image'] = fluid.data(
  212. dtype='float32', shape=[None, 3, None, None], name='image')
  213. if self.mode == 'train':
  214. inputs['label'] = fluid.data(
  215. dtype='int32', shape=[None, 1, None, None], name='label')
  216. elif self.mode == 'eval':
  217. inputs['label'] = fluid.data(
  218. dtype='int32', shape=[None, 1, None, None], name='label')
  219. return inputs
  220. def build_net(self, inputs):
  221. # 在两类分割情况下,当loss函数选择dice_loss或bce_loss的时候,最后logit输出通道数设置为1
  222. if self.use_dice_loss or self.use_bce_loss:
  223. self.num_classes = 1
  224. image = inputs['image']
  225. encode_data, short_cuts = self._encode(image)
  226. decode_data = self._decode(encode_data, short_cuts)
  227. logit = self._get_logit(decode_data, self.num_classes)
  228. if self.num_classes == 1:
  229. out = sigmoid_to_softmax(logit)
  230. out = fluid.layers.transpose(out, [0, 2, 3, 1])
  231. else:
  232. out = fluid.layers.transpose(logit, [0, 2, 3, 1])
  233. pred = fluid.layers.argmax(out, axis=3)
  234. pred = fluid.layers.unsqueeze(pred, axes=[3])
  235. if self.mode == 'train':
  236. label = inputs['label']
  237. mask = label != self.ignore_index
  238. return self._get_loss(logit, label, mask)
  239. elif self.mode == 'eval':
  240. label = inputs['label']
  241. mask = label != self.ignore_index
  242. loss = self._get_loss(logit, label, mask)
  243. return loss, pred, label, mask
  244. else:
  245. if self.num_classes == 1:
  246. logit = sigmoid_to_softmax(logit)
  247. else:
  248. logit = fluid.layers.softmax(logit, axis=1)
  249. return pred, logit