"""Map the model's native/explicit grounding slots to ordered input images."""


def image_mask_slots(pred_masks, num_images):
    if pred_masks is None:
        return [None] * num_images
    if len(pred_masks) == num_images + 1:
        pred_masks = pred_masks[1:]
    if len(pred_masks) != num_images:
        raise ValueError('Grounding output does not match the number of input images.')
    return [slot[0] if slot else None for slot in pred_masks]