Freq weight IoU

from PIL import Image
import numpy as np
import tensorflow as tf


def fre_iou(labels,
            predictions,
            num_classes,
            weights=None,
            metrics_collections=None,
            updates_collections=None,
            name=None):
    """Calculate per-step mean Intersection-Over-Union (mIOU).

    Mean Intersection-Over-Union is a common evaluation metric for
    semantic image segmentation, which first computes the IOU for each
    semantic class and then computes the average over classes.
    IOU is defined as follows:
      IOU = true_positive / (true_positive + false_positive + false_negative).
    The predictions are accumulated in a confusion matrix, weighted by `weights`,
    and mIOU is then calculated from it.

    For estimation of the metric over a stream of data, the function creates an
    `update_op` operation that updates these variables and returns the `mean_iou`.

    If `weights` is `None`, weights default to 1. Use weights of 0 to mask values.

    Args:
      labels: A `Tensor` of ground truth labels with shape [batch size] and of
        type `int32` or `int64`. The tensor will be flattened if its rank > 1.
      predictions: A `Tensor` of prediction results for semantic labels, whose
        shape is [batch size] and type `int32` or `int64`. The tensor will be
        flattened if its rank > 1.
      num_classes: The possible number of labels the prediction task can
        have. This value must be provided, since a confusion matrix of
        dimension = [num_classes, num_classes] will be allocated.
      weights: Optional `Tensor` whose rank is either 0, or the same rank as
        `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
        be either `1`, or the same as the corresponding `labels` dimension).
      metrics_collections: An optional list of collections that `mean_iou`
        should be added to.
      updates_collections: An optional list of collections `update_op` should be
        added to.
      name: An optional variable_scope name.

    Returns:
      mean_iou: A `Tensor` representing the mean intersection-over-union.
      update_op: An operation that increments the confusion matrix.

    Raises:
      ValueError: If `predictions` and `labels` have mismatched shapes, or if
        `weights` is not `None` and its shape doesn't match `predictions`, or if
        either `metrics_collections` or `updates_collections` are not a list or
        tuple.
    """
    with tf.variable_scope(
            name, 'freq_iou', (predictions, labels, weights)):
        # Check if shape is compatible.
        predictions.get_shape().assert_is_compatible_with(labels.get_shape())
        total_cm = tf.confusion_matrix(labels, predictions=predictions, num_classes=num_classes, weights=weights)
        total_cm, update_op = streaming_confusion_matrix(labels, predictions,
                                                         num_classes, weights)

        def compute_freq_iou(name):
            """Compute the mean intersection-over-union via the confusion matrix."""
            sum_over_row = tf.to_float(tf.reduce_sum(total_cm, 0))
            sum_over_col = tf.to_float(tf.reduce_sum(total_cm, 1))
            cm_diag = tf.to_float(tf.diag_part(total_cm))
            denominator = sum_over_row + sum_over_col - cm_diag

            # If the value of the denominator is 0, set it to 1 to avoid
            # zero division.
            denominator = tf.where(
                tf.greater(denominator, 0),
                denominator,
                tf.ones_like(denominator))

            iou = tf.div(cm_diag, denominator)
            # return tf.reduce_mean(iou, name=name)
            # Freq weight IoU
            fiou = tf.div(tf.reduce_sum(tf.multiply(sum_over_row, iou)), tf.to_float(tf.reduce_sum(total_cm)),
                          name=name)
            return fiou

        freq_iou_v = compute_freq_iou('freq_iou')

        if metrics_collections:
            tf.ops.add_to_collections(metrics_collections, freq_iou_v)

        if updates_collections:
            tf.ops.add_to_collections(updates_collections, update_op)

        return freq_iou_v, update_op


def streaming_confusion_matrix(labels, predictions, num_classes, weights=None):
    """Calculate a streaming confusion matrix.

    Calculates a confusion matrix. For estimation over a stream of data,
    the function creates an  `update_op` operation.

    Args:
      labels: A `Tensor` of ground truth labels with shape [batch size] and of
        type `int32` or `int64`. The tensor will be flattened if its rank > 1.
      predictions: A `Tensor` of prediction results for semantic labels, whose
        shape is [batch size] and type `int32` or `int64`. The tensor will be
        flattened if its rank > 1.
      num_classes: The possible number of labels the prediction task can
        have. This value must be provided, since a confusion matrix of
        dimension = [num_classes, num_classes] will be allocated.
      weights: Optional `Tensor` whose rank is either 0, or the same rank as
        `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
        be either `1`, or the same as the corresponding `labels` dimension).

    Returns:
      total_cm: A `Tensor` representing the confusion matrix.
      update_op: An operation that increments the confusion matrix.
    """
    # Local variable to accumulate the predictions in the confusion matrix.
    cm_dtype = tf.int64 if weights is not None else tf.float64
    total_cm = tf.get_local_variable(
        'total_confusion_matrix',
        shape=[num_classes, num_classes],
        dtype=cm_dtype)

    # Cast the type to int64 required by confusion_matrix_ops.
    predictions = tf.to_int64(predictions)
    labels = tf.to_int64(labels)
    num_classes = tf.to_int64(num_classes)

    # Flatten the input if its rank > 1.
    if predictions.get_shape().ndims > 1:
        predictions = tf.array_ops.reshape(predictions, [-1])

    if labels.get_shape().ndims > 1:
        labels = tf.array_ops.reshape(labels, [-1])

    if (weights is not None) and (weights.get_shape().ndims > 1):
        weights = tf.array_ops.reshape(weights, [-1])

    # Accumulate the prediction to current confusion matrix.
    current_cm = tf.confusion_matrix(
        labels, predictions, num_classes, weights=weights, dtype=cm_dtype)
    # tf.assign_add()
    update_op = tf.assign_add(total_cm, current_cm)
    return total_cm, update_op