Calculating IOU for semantic segmentation


I want to calculate the mean Intersection over Union (mIoU) of my predicted vs ground truth semantic segmentation labels.

My prediction is of shape [B, 1, H, W] where B is the batch size, H is the height and W is the width. Each pixel has an integer that represents the class (eg: 1-30). The labels are the same shape, including the integer class.

I would like a way to calculate and report the mIoU based on these two tensors. Is there a library or implementation somewhere that I can refer to?

Thank you!