Can anyone please check my code for calculating IoU (0 & 1 classes) and mIoU scores during the validation? Is this correct?
train_losses = []
val_losses = []
# Function to compute IoU
def compute_iou_binary_torch(y_true, y_pred):
"""
Computes IoU for binary classification directly from PyTorch tensors.
Args:
y_true (tensor): Ground truth binary tensor (0s and 1s).
y_pred (tensor): Predicted binary tensor (0s and 1s).
Returns:
IoU (list): Intersection over Union for each class.
mIoU (float): Mean IoU across classes.
"""
# Convert tensors to NumPy arrays
y_true = y_true.cpu().numpy().flatten()
y_pred = y_pred.cpu().numpy().flatten()
# Compute confusion matrix
current = confusion_matrix(y_true, y_pred, labels=[1, 0])
intersection = np.diag(current)
ground_truth_set = current.sum(axis=1)
predicted_set = current.sum(axis=0)
union = (ground_truth_set + predicted_set - intersection).astype(np.float32)
with np.errstate(divide='ignore', invalid='ignore'):
IoU = np.divide(intersection, union)
IoU[np.isnan(IoU)] = 0 # Replace NaN values with 0
mIoU = np.nanmean(IoU)
return IoU, mIoU
# Training and validation loop
for epoch in range(n_eps):
model.train()
train_loss = 0.0
# Training loop
for images, masks in tqdm(train_loader):
images, masks = images.to(device), masks.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, masks)
loss.backward()
optimizer.step()
train_loss += loss.item()
avg_train_loss = train_loss / len(train_loader)
train_losses.append(avg_train_loss)
print(f"Epoch [{epoch+1}/{n_eps}], Train Loss: {avg_train_loss:.4f}")
model.eval()
val_loss = 0.0
iou_scores = [] # Reset IoU scores at the start of the validation loop
class_iou_scores = [] # Reset class-wise IoU scores at the start of the validation loop
# Validation loop
with torch.no_grad():
for images, masks in val_loader:
images, masks = images.to(device), masks.to(device)
outputs = model(images)
# Calculate loss on raw outputs
val_loss += criterion(outputs, masks).item()
# Apply sigmoid and thresholding for binary prediction
predictions = (outputs.sigmoid() > 0.5).float()
# Compute IoU for the current batch
batch_iou, batch_miou = compute_iou_binary_torch(masks, predictions)
class_iou_scores.append(batch_iou)
iou_scores.append(batch_miou)
avg_val_loss = val_loss / len(val_loader)
avg_iou = np.mean(iou_scores) # Compute average IoU across all batches
avg_class_iou = np.mean(class_iou_scores, axis=0) # Compute average IoU for each class
val_losses.append(avg_val_loss)
print(f"Epoch [{epoch+1}/{n_eps}], Val Loss: {avg_val_loss:.4f}")
print(f"Epoch [{epoch+1}/{n_eps}], Class-wise IoU: {avg_class_iou}")
print(f"Epoch [{epoch+1}/{n_eps}], Mean IoU: {avg_iou:.4f}")