Repository: Project-MONAI/MONAI
File: monai/metrics/generalized_dice.py (Python)

Code before:
    if not include_background:
        y_pred, y = ignore_background(y_pred=y_pred, y=y)

    # Reducing only spatial dimensions (not batch nor channels), compute the intersection and non-weighted denominator
    reduce_axis = list(range(2, y_pred.dim()))
    intersection = torch.sum(y * y_pred, dim=reduce_axis)
    y_o = torch.sum(y, dim=reduce_axis)
    y_pred_o = torch.sum(y_pred, dim=reduce_axis)
    denominator = y_o + y_pred_o

    # Set the class weights
    weight_type = look_up_option(weight_type, Weight)
    if weight_type == Weight.SIMPLE:
        w = torch.reciprocal(y_o.float())
    elif weight_type == Weight.SQUARE:
        w = torch.reciprocal(y_o.float() * y_o.float())
    else:
        w = torch.ones_like(y_o.float())

    # Replace infinite values for non-appearing classes by the maximum weight
    for b in w:
        infs = torch.isinf(b)
        b[infs] = 0
        b[infs] = torch.max(b)

    # Compute the weighted numerator and denominator, summing along the class axis when sum_over_labels is True
    if sum_over_classes:
        numer = 2.0 * (intersection * w).sum(dim=1, keepdim=True)
        denom = (denominator * w).sum(dim=1, keepdim=True)
        y_pred_o = y_pred_o.sum(dim=-1, keepdim=True)
    else:
        numer = 2.0 * (intersection * w)
        denom = denominator * w
        y_pred_o = y_pred_o

    # Compute the score
    generalized_dice_score = numer / denom

    # Handle zero division. Where denom == 0 and the prediction volume is 0, score is 1.
    # Where denom == 0 but the prediction volume is not 0, score is 0
    denom_zeros = denom == 0
    generalized_dice_score[denom_zeros] = torch.where(
        (y_pred_o == 0)[denom_zeros],
        torch.tensor(1.0, device=generalized_dice_score.device),
        torch.tensor(0.0, device=generalized_dice_score.device),
    )

    return generalized_dice_score


Review comment (suggestion):
```suggestion
    # Compute the weighted numerator and denominator, summing along the class axis when sum_over_classes is True
```

Code after:
    if not include_background:
        y_pred, y = ignore_background(y_pred=y_pred, y=y)

    # Reducing only spatial dimensions (not batch nor channels), compute the intersection and non-weighted denominator
    reduce_axis = list(range(2, y_pred.dim()))
    intersection = torch.sum(y * y_pred, dim=reduce_axis)
    y_o = torch.sum(y, dim=reduce_axis)
    y_pred_o = torch.sum(y_pred, dim=reduce_axis)
    denominator = y_o + y_pred_o

    # Set the class weights
    weight_type = look_up_option(weight_type, Weight)
    if weight_type == Weight.SIMPLE:
        w = torch.reciprocal(y_o.float())
    elif weight_type == Weight.SQUARE:
        w = torch.reciprocal(y_o.float() * y_o.float())
    else:
        w = torch.ones_like(y_o.float())

    # Replace infinite values for non-appearing classes by the maximum weight
    for b in w:
        infs = torch.isinf(b)
        b[infs] = 0
        b[infs] = torch.max(b)

    # Compute the weighted numerator and denominator, summing along the class axis when sum_over_classes is True
    if sum_over_classes:
        numer = 2.0 * (intersection * w).sum(dim=1, keepdim=True)
        denom = (denominator * w).sum(dim=1, keepdim=True)
        y_pred_o = y_pred_o.sum(dim=-1, keepdim=True)
    else:
        numer = 2.0 * (intersection * w)
        denom = denominator * w
        y_pred_o = y_pred_o

    # Compute the score
    generalized_dice_score = numer / denom

    # Handle zero division. Where denom == 0 and the prediction volume is 0, score is 1.
    # Where denom == 0 but the prediction volume is not 0, score is 0
    denom_zeros = denom == 0
    generalized_dice_score[denom_zeros] = torch.where(
        (y_pred_o == 0)[denom_zeros],
        torch.tensor(1.0, device=generalized_dice_score.device),
        torch.tensor(0.0, device=generalized_dice_score.device),
    )

    return generalized_dice_score