import torch
import torchmetrics
import pandas as pd

num_classes = 20
batch_size = 16
train_metrics = get_metric_collection(num_classes)
dice_score_default = torchmetrics.segmentation.DiceScore(
    input_format="index",
    num_classes=num_classes,
)
dice_score_no_bg = torchmetrics.segmentation.DiceScore(
    input_format="index",
    num_classes=num_classes,
    include_background=False,
)
dice_score_macro = torchmetrics.segmentation.DiceScore(
    input_format="index",
    num_classes=num_classes,
    average="macro",
)
dice_score_realistic = torchmetrics.segmentation.DiceScore(
    input_format="index",
    num_classes=num_classes,
    include_background=False,
    average="macro",
)

# Create example output and target tensors, where the background is 75% of the output
example_output = torch.randint(0, num_classes, (batch_size, 10, 10))
output_background_mask = torch.rand(example_output.shape) < 0.75
example_output[output_background_mask] = 0

# Create example target tensor, where the background is 75% of the target
example_target = torch.randint(0, num_classes, (batch_size, 10, 10))
target_background_mask = torch.rand(example_target.shape) < 0.75
example_target[target_background_mask] = 0

dice_score_default.update(example_output, example_target)
dice_score_no_bg.update(example_output, example_target)
dice_score_macro.update(example_output, example_target)
dice_score_realistic.update(example_output, example_target)

scores = {
    "include_background": ["True", "False"],
    "average='micro'": [
        dice_score_default.compute().item(),
        dice_score_no_bg.compute().item(),
    ],
    "average='macro'": [
        dice_score_macro.compute().item(),
        dice_score_realistic.compute().item(),
    ],
}
scores_df = pd.DataFrame(scores)
print(scores_df)
#   include_background  average='micro'  average='macro'
# 0               True         0.575000         0.042818
# 1              False         0.007878         0.005482