add psenet model file

nets/model.py

@@ -0,0 +1,152 @@
+#-*- coding:utf-8 -*-
+import tensorflow as tf
+from utils.utils_tool import logger
+
+from tensorflow.contrib import slim
+
+tf.app.flags.DEFINE_integer('text_scale', 512, '')
+
+from nets.resnet import resnet_v1
+
+FLAGS = tf.app.flags.FLAGS
+
+#TODO:使用最近领还是线性插值，文章没有介绍
+def unpool(inputs, rate):
+    return tf.image.resize_bilinear(inputs, size=[tf.shape(inputs)[1]*rate,  tf.shape(inputs)[2]*rate])
+
+def mean_image_subtraction(images, means=[123.68, 116.78, 103.94]):
+    '''
+    image normalization
+    :param images:
+    :param means:
+    :return:
+    '''
+    num_channels = images.get_shape().as_list()[-1]
+    if len(means) != num_channels:
+      raise ValueError('len(means) must match the number of channels')
+    channels = tf.split(axis=3, num_or_size_splits=num_channels, value=images)
+    for i in range(num_channels):
+        channels[i] -= means[i]
+    return tf.concat(axis=3, values=channels)
+
+def build_feature_pyramid(C, weight_decay):
+
+    '''
+    reference: https://github.com/CharlesShang/FastMaskRCNN
+    build P2, P3, P4, P5
+    :return: multi-scale feature map
+    '''
+
+    feature_pyramid = {}
+    with tf.variable_scope('build_feature_pyramid'):
+        with slim.arg_scope([slim.conv2d], weights_regularizer=slim.l2_regularizer(weight_decay)):
+            feature_pyramid['P5'] = slim.conv2d(C['C5'],
+                                                num_outputs=256,
+                                                kernel_size=[1, 1],
+                                                stride=1,
+                                                scope='build_P5')
+
+            feature_pyramid['P6'] = slim.max_pool2d(feature_pyramid['P5'],
+                                                    kernel_size=[2, 2], stride=2, scope='build_P6')
+            # P6 is down sample of P5
+
+            for layer in range(4, 1, -1):
+                p, c = feature_pyramid['P' + str(layer + 1)], C['C' + str(layer)]
+                up_sample_shape = tf.shape(c)
+                up_sample = tf.image.resize_nearest_neighbor(p, [up_sample_shape[1], up_sample_shape[2]],
+                                                             name='build_P%d/up_sample_nearest_neighbor' % layer)
+
+                c = slim.conv2d(c, num_outputs=256, kernel_size=[1, 1], stride=1,
+                                scope='build_P%d/reduce_dimension' % layer)
+                p = up_sample + c
+                p = slim.conv2d(p, 256, kernel_size=[3, 3], stride=1,
+                                padding='SAME', scope='build_P%d/avoid_aliasing' % layer)
+                feature_pyramid['P' + str(layer)] = p
+    return feature_pyramid
+
+def model(images, outputs = 6, weight_decay=1e-5, is_training=True):
+    '''
+    define the model, we use slim's implemention of resnet
+    '''
+    images = mean_image_subtraction(images)
+
+    with slim.arg_scope(resnet_v1.resnet_arg_scope(weight_decay=weight_decay)):
+        logits, end_points = resnet_v1.resnet_v1_50(images, is_training=is_training, scope='resnet_v1_50')
+
+    #no non-linearities in FPN article
+    feature_pyramid = build_feature_pyramid(end_points, weight_decay=weight_decay)
+    #unpool sample P
+    P_concat = []
+    for i in range(3, 1, -1):
+        P_concat.append(unpool(feature_pyramid['P'+str(i+2)], 2**i))
+    P_concat.append(feature_pyramid['P2'])
+    #F = C(P2,P3,P4,P5)
+    F = tf.concat(P_concat, axis=-1)
+
+    #reduce to 256 channels
+    with tf.variable_scope('feature_results'):
+        batch_norm_params = {
+            'decay': 0.997,
+            'epsilon': 1e-5,
+            'scale': True,
+            'is_training': is_training
+        }
+        with slim.arg_scope([slim.conv2d],
+                            activation_fn=tf.nn.relu,
+                            normalizer_fn=slim.batch_norm,
+                            normalizer_params=batch_norm_params,
+                            weights_regularizer=slim.l2_regularizer(weight_decay)):
+            F = slim.conv2d(F, 256, 3)
+        with slim.arg_scope([slim.conv2d],
+                            weights_regularizer=slim.l2_regularizer(weight_decay),
+                            activation_fn=None):
+            S = slim.conv2d(F, outputs, 1)
+    up_S = unpool(S, 4)
+    seg_S_pred = tf.nn.sigmoid(up_S)
+    shape_res_S = tf.shape(seg_S_pred)
+    logger.debug(shape_res_S)
+
+    return seg_S_pred
+
+def dice_coefficient(y_true_cls, y_pred_cls,
+                     training_mask):
+    '''
+    dice loss
+    :param y_true_cls: ground truth
+    :param y_pred_cls: predict
+    :param training_mask:
+    :return:
+    '''
+    eps = 1e-5
+    intersection = tf.reduce_sum(y_true_cls * y_pred_cls * training_mask)
+    union = tf.reduce_sum(y_true_cls * training_mask) + tf.reduce_sum(y_pred_cls * training_mask) + eps
+    dice = 2 * intersection / union
+    loss = 1. - dice
+    # tf.summary.scalar('classification_dice_loss', loss)
+    return dice, loss
+
+def loss(y_true_cls, y_pred_cls,
+         training_mask):
+    g1, g2, g3, g4, g5, g6 = tf.split(value=y_true_cls, num_or_size_splits=6, axis=3)
+    s1, s2, s3, s4, s5, s6 = tf.split(value=y_pred_cls, num_or_size_splits=6, axis=3)
+    Gn = [g1, g2, g3, g4, g5, g6]
+    Sn = [s1, s2, s3, s4, s5, s6]
+    _, Lc = dice_coefficient(Gn[5], Sn[5], training_mask=training_mask)
+    tf.summary.scalar('Lc_loss', Lc)
+
+    one = tf.ones_like(Sn[5])
+    zero = tf.zeros_like(Sn[5])
+    W = tf.where(Sn[5] >= 0.5, x=one, y=zero)
+    D = 0
+    for i in range(5):
+        di, _ = dice_coefficient(Gn[i]*W, Sn[i]*W, training_mask=training_mask)
+        D += di
+    Ls = 1-D/5.
+    tf.summary.scalar('Ls_loss', Ls)
+    lambda_ = 0.7
+    L = lambda_*Lc + (1-lambda_)*Ls
+    return L
+
+
+
+