Correct way to calculate IoU for each class

I am doing multi class segmentation and I want to know what is the correct way for calculating and displaying iou for each class during the validation of the data. Any help will be appreciated. Thank you

Below is my code:

Validation

    model.eval()
    running_val_loss = 0.0
    with torch.no_grad():
        for images, masks in val_loader:  # Replace val_loader with your DataLoader
            images, masks = images.to(device), masks.to(device)

            outputs = model(images)

            # Calculate validation loss
            val_loss = criterion(outputs, masks)
            running_val_loss += val_loss.item()

            # Call the function to calculate mIoU and IoU metrics
            miou = calculate_iou_scores(outputs, masks, num_classes=5)


    avg_val_loss = running_val_loss / len(val_loader)
    val_losses.append(avg_val_loss)

    # Print validation results
    print(f'Epoch [{epoch+1}/{n_eps}], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')

calculate_iou_score function

def calculate_iou_scores(outputs, targets, num_classes):
    # Convert logits to class predictions
    outputs = torch.argmax(outputs, dim=1)

    # Initialize IoU metric
    iou = torchmetrics.JaccardIndex(num_classes=num_classes, average='macro', task='multiclass').to(device)


    # Calculate IoU and mIoU
    iou.update(outputs, targets)
    iou_score = iou.compute()

    return iou_score.item()

Something like this may work:

# Val
model.eval()
running_val_loss = 0.0
class_iou_sums = [0] * num_classes
with torch.no_grad():
    for images, masks in val_loader:
        outputs = model(images)
        val_loss = criterion(outputs, masks)
        running_val_loss += val_loss.item()
        iou_scores, batch_miou = calculate_iou_scores(outputs, masks, num_classes=num_classes)
        for i, iou in enumerate(iou_scores):
            class_iou_sums[i] += iou

avg_val_loss = running_val_loss / len(val_loader)
class_ious = [iou_sum / len(val_loader) for iou_sum in class_iou_sums]
miou = sum(class_ious) / num_classes

# Iterative Results
print(f'Epoch [1/1], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')
print('IoU for each class:')
for i, iou in enumerate(class_ious):
    print(f'Class {i}: {iou:.4f}')

this code is not working…!

It is on my machine (haha)

Also, do you mean not working as in logically or its throwing some error?

# pip install torchmetrics
import torch
import torchmetrics

# Dummy values
num_classes = 5
num_batches = 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Dummy model (just for testing)
class DummyModel(torch.nn.Module):
    def forward(self, x):
        return torch.randn(x.size(0), num_classes, 16, 16)

model = DummyModel().to(device)

criterion = torch.nn.CrossEntropyLoss()

# Dummy validation loader
class DummyLoader:
    def __init__(self, num_batches):
        self.num_batches = num_batches

    def __iter__(self):
        for _ in range(self.num_batches):
            images = torch.randn(4, 3, 16, 16).to(device)
            masks = torch.randint(0, num_classes, (4, 16, 16)).to(device)
            yield images, masks

    def __len__(self):
        return self.num_batches

val_loader = DummyLoader(num_batches)

def calculate_iou_scores(outputs, targets, num_classes):
    outputs = torch.argmax(outputs, dim=1)
    iou_metrics = [torchmetrics.JaccardIndex(num_classes=num_classes, average='none', task='multiclass').to(device)
                   for _ in range(num_classes)]
    iou_scores = []
    for class_idx in range(num_classes):
        iou_metrics[class_idx].update(outputs, targets)
        iou_score = iou_metrics[class_idx].compute()
        iou_scores.append(iou_score[class_idx].item())
    miou = sum(iou_scores) / num_classes
    return iou_scores, miou

# Val loop
model.eval()
running_val_loss = 0.0
class_iou_sums = [0] * num_classes
with torch.no_grad():
    for images, masks in val_loader:
        outputs = model(images)
        val_loss = criterion(outputs, masks)
        running_val_loss += val_loss.item()
        iou_scores, batch_miou = calculate_iou_scores(outputs, masks, num_classes=num_classes)
        for i, iou in enumerate(iou_scores):
            class_iou_sums[i] += iou

avg_val_loss = running_val_loss / len(val_loader)
class_ious = [iou_sum / len(val_loader) for iou_sum in class_iou_sums]
miou = sum(class_ious) / num_classes

# Print validation results
print(f'Epoch [1/1], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')
print('IoU for each class:')
for i, iou in enumerate(class_ious):
    print(f'Class {i}: {iou:.4f}')

Maybe try running this.

1 Like

@Soumya_Kundu thanks man it is working
previously there was issue in the calculate_iou_scores function