import torch


def weighted_jaccard_index(target, predictions):
    """
    Semantic segmentation metric. Calculates mean intersection over union over all classes, weighted by class.
    Works for batches of data.
    """

    class_counts = torch.sum(target, dim=(-2, -1), dtype=torch.int32)
    n_pixels = torch.sum(class_counts, dim=(-1))[0]
    class_frequencies = class_counts / n_pixels

    intersection = torch.logical_and(target, predictions)
    union = torch.logical_or(target, predictions)

    intersection_sums = torch.sum(intersection, dim=(-2, -1))
    union_sums = torch.sum(union, dim=(-2,-1))

    # union_sums will contain 0's if a class is not present in an image, which will give division by zero
    iou_scores_classwise = intersection_sums / (union_sums + 0.00000000001)

    iou_scores_weighted = torch.sum(iou_scores_classwise * class_frequencies, dim=(-1))

    iou_score_batch_weighted_mean = torch.mean(iou_scores_weighted)

    return iou_score_batch_weighted_mean.numpy()