src/transformers/models/oneformer/convert_to_hf_oneformer.py/post_process_sem_seg_output
def post_process_sem_seg_output(outputs: OneFormerForUniversalSegmentationOutput, target_size: Tuple[int, int]):
    # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
    class_queries_logits = outputs.class_queries_logits
    # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_queries_logits = outputs.masks_queries_logits
    if target_size is not None:
        masks_queries_logits = torch.nn.functional.interpolate(
            masks_queries_logits,
            size=target_size,
            mode="bilinear",
            align_corners=False,
        )
    # remove the null class `[..., :-1]`
    masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
    # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_probs = masks_queries_logits.sigmoid()
    # now we want to sum over the queries,
    # $ out_{c,h,w} =  \sum_q p_{q,c} * m_{q,h,w} $
    # where $ softmax(p) \in R^{q, c} $ is the mask classes
    # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
    # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
    segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

    return segmentation