import torch


def blend_seg_pred(img, seg, alpha=0.5):
    """ Blend images with their corresponding segmentation prediction.

    Args:
        img (torch.Tensor): A batch of image tensors of shape (B, 3, H, W) where B is the batch size,
            H is the images height and W is the images width
        seg (torch.Tensor): A batch of segmentation predictions of shape (B, C, H, W) where B is the batch size,
            C is the number of segmentation classes, H is the images height and W is the images width
        alpha: alpha (float): Opacity value for the segmentation in the range [0, 1] where 0 is completely transparent
            and 1 is completely opaque

    Returns:
        torch.Tensor: The blended image.
    """
    pred = seg.argmax(1)
    pred = pred.view(pred.shape[0], 1, pred.shape[1], pred.shape[2]).repeat(1, 3, 1, 1)
    blend = img

    # For each segmentation class except the background (label 0)
    for i in range(1, seg.shape[1]):
        color_mask = -torch.ones_like(img)
        color_mask[:, -i, :, :] = 1
        alpha_mask = 1 - (pred == i).float() * alpha
        blend = blend * alpha_mask + color_mask * (1 - alpha_mask)

    return blend