import torch


def multiclass_dice_threshold(logits, targets, threshold=0.5, iou=False, eps=1e-8):
    """
    Dice coefficient metric for multiclass binary target. 
    If iou=True, returns iou metric, classic for segmentation problems.
    """

    n = targets.shape[0]  # Batch size of 4

    # Flatten logits and targets
    logits = logits.view(n, -1)
    targets = targets.view(n, -1).float()

    # Convert logits to probabilities
    probs = torch.sigmoid(logits)

    preds = probs
    preds[preds >= threshold] = 1
    preds[preds < threshold] = 0

    intersect = (preds * targets).sum(dim=1).float()
    union = (preds + targets).sum(dim=1).float()

    if not iou:
        l = 2. * intersect / union
    else:
        l = intersect / (union - intersect + eps)

    # The Dice coefficient is defined to be 1 when both X and Y are empty.
    # That said, we'd get a divide-by-zero-exception if union was 0 anyways...
    l[union == 0.] = 1.
    return l.mean()