import torch


def instance_masks_to_semseg_mask(instance_masks, category_labels):
    """
    Converts a tensor containing instance masks to a semantic segmentation mask.
    :param instance_masks: tensor(N, T, H, W)  (N = number of instances)
    :param category_labels: tensor(N) containing semantic category label for each instance.
    :return: semantic mask as tensor(T, H, W] with pixel values containing class labels
    """
    assert len(category_labels) == instance_masks.shape[0], \
        "Number of instances do not match: {}, {}".format(len(category_labels), len(instance_masks))
    semseg_masks = instance_masks.long()

    for i, label in enumerate(category_labels):
        semseg_masks[i] = torch.where(instance_masks[i], label, semseg_masks[i])

    # for pixels with differing labels, assign to the category with higher ID number (arbitrary criterion)
    return semseg_masks.max(dim=0)[0]