[人工智能]maskrcnn边缘检测

def mrcnn_edge_loss_graph(target_masks, target_class_ids, pred_masks):
    """Edge L2 loss for mask edge head

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of Value 0 or 1(boolean?). Use zero padding to fill array
    target_class_ids: [batch, num_rois]. Integer class IDs. Zeros padded.
    pred_masks: [batch, proposal, height, width, num_classes] float32 tensor
                with value from 0 to 1(soft mask)(more information)
    """
    # Reshape for simplicity. Merge first two dimensions into one
    # 即将batch 和 num_rois 合并为一项
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. (正的ROI是相对BG而言吗)
    # And only the class specific mask of each ROI
    # tf.where 获得索引值
    # tf.gather 根据索引值从张量中获得元素构成新张量Tensor
    # tf.cast 类型转换
    # tf.stack
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # shape: [batch * rois, height, width, 1]
    y_true = tf.expand_dims(y_true, -1)
    y_pred = tf.expand_dims(y_pred, -1)

    y_true = 255 * y_true
    y_pred = 255 * y_pred

    # shape: [3, 3, 1, 2]
    sobel_kernel = tf.constant([[[[1, 1]], [[0, 2]], [[-1, 1]]],
                                [[[2, 0]], [[0, 0]], [[-2, 0]]],
                                [[[1, -1]], [[0, -2]], [[-1, -1]]]], dtype=tf.float32)

    # Conv2D with kernel
    edge_true = tf.nn.conv2d(y_true, sobel_kernel, strides=[1, 1, 1, 1], padding="SAME")
    edge_pred = tf.nn.conv2d(y_pred, sobel_kernel, strides=[1, 1, 1, 1], padding="SAME")

    # abs and clip
    edge_true = tf.clip_by_value(abs(edge_true), 0, 255)
    edge_pred = tf.clip_by_value(abs(edge_pred), 0, 255)

    # Mean Square Error(MSE) Loss
    return tf.reduce_mean(tf.square(edge_true / 255. - edge_pred / 255.))