def process_panoptic_prediction()

in seamseg/utils/coco_ap.py [0:0]

• 24 lines of code
• 6 McCabe index (conditional complexity)

def process_panoptic_prediction(panoptic_pred, num_stuff, idx, img_size, original_size):
    # Extract panoptic prediction
    msk_pred, cat_pred, obj_pred, iscrowd_pred = panoptic_pred

    bbx_pred = extract_boxes(msk_pred, cat_pred.numel())

    # Convert bbx and redo clamping
    bbx_pred[:, [0, 2]] = (bbx_pred[:, [0, 2]] / img_size[0] * original_size[0]).clamp(min=0, max=original_size[0])
    bbx_pred[:, [1, 3]] = (bbx_pred[:, [1, 3]] / img_size[1] * original_size[1]).clamp(min=0, max=original_size[1])
    bbx_pred_size = bbx_pred[:, 2:] - bbx_pred[:, :2]

    outs = []
    for i, (obj_i, cat_i, bbx_i, iscrowd_i, bbx_size_i) in enumerate(zip(
            obj_pred, cat_pred, bbx_pred, iscrowd_pred, bbx_pred_size)):
        if iscrowd_i.item() == 1 or cat_i.item() < num_stuff or cat_i.item() == 255:
            continue
        out = dict(image_id=idx, category_id=int(cat_i.item()), score=float(obj_i.item()))

        out["bbox"] = [
            float(bbx_i[1].item()),
            float(bbx_i[0].item()),
            float(bbx_size_i[1].item()),
            float(bbx_size_i[0].item()),
        ]

        segmentation = msk_pred == i
        segmentation = Image.fromarray(segmentation.numpy()).resize(original_size[::-1], Image.NEAREST)
        out["segmentation"] = mask_encode(np.asfortranarray(np.array(segmentation)))
        out["segmentation"]["counts"] = str(out["segmentation"]["counts"], "utf-8")

        outs.append(out)

    return outs