import datasets
from torchmetrics.classification import JaccardIndex
from torchvision.transforms.v2.functional import pil_to_tensor

TAXONOMIES = {
    'cityscapes': {
        'num_classes': 19,
        'label_col': 'label',
        'gt_offset': 0
    },
    'ade20k': {
        'num_classes': 150,
        'label_col': 'annotation',
        'gt_offset': -1
    },
}

def evaluate_miou(model_fn, taxonomy, dataset):
    cfg = taxonomy if isinstance(taxonomy, dict) else TAXONOMIES[taxonomy]
    ds = datasets.load_dataset(dataset, split='validation')
    metric = JaccardIndex(
        task="multiclass",
        num_classes=cfg['num_classes'],
        average="macro",
        ignore_index=255)
    for sample in ds:
        sample = model_fn(sample)
        pred = pil_to_tensor(sample['pred'])
        gt = pil_to_tensor(sample[cfg['label_col']]) + cfg['gt_offset']
        metric.update(pred, gt)
    return metric.compute().item()