def preparePredImage()

in maskrcnn_benchmark/data/datasets/evaluation/cityscapes/eval_instances.py [0:0]
49 lines of code
18 McCabe index (conditional complexity)

def preparePredImage(dataset, predictions, idx):
    perImagePredictions = predictions[idx]

    # A list will hold statistics and meta-data about the image
    perImageInstances = []

    # maskTensor represents binary masks of all predicted instance segmentations
    # if present
    maskTensor = [None] * len(perImagePredictions)

    # No predictions for this image
    if len(perImagePredictions) == 0:
        return perImageInstances, maskTensor

    # Resize to original image size
    imgInfo = dataset.get_img_info(idx)
    origSize = imgInfo["width"], imgInfo["height"]
    if perImagePredictions.size != origSize:
        perImagePredictions = perImagePredictions.resize(size=origSize)

    # Bounding boxes and areas
    perImagePredictions = perImagePredictions.convert("xyxy")
    bbs = perImagePredictions.bbox.long()
    xmins, ymins, xmaxs, ymaxs = bbs[:, 0], bbs[:, 1], bbs[:, 2], bbs[:, 3]
    boxAreas = ((xmaxs - xmins) * (ymaxs - ymins)).tolist()
    bbs = bbs.tolist()

    # object label and "Objectness" score for each prediction
    labels = perImagePredictions.get_field("labels").tolist()
    scores = perImagePredictions.get_field("scores").tolist()

    # Get the mask for each instance in a contiguous array
    if "mask" in perImagePredictions.fields():
        maskTensor = perImagePredictions.get_field("mask")

        # sanity checks
        assert len(perImagePredictions) == len(maskTensor), (
            "number of masks (%d) do not match the number of boxes (%d)"
            % (len(perImagePredictions), len(maskTensor))
        )

        maskTensor = maskTensor.float()
        # We assume that the maskTensors are coming right out of the maskRCNN
        # having values between [0, 1] inclusive
        #
        # assert maskTensor.min() >= 0.0 and maskTensor.max() <= 1.0, [
        #     maskTensor.max(),
        #     maskTensor.min(),
        # ]

        # Project masks to the boxes
        # TODO: Issue #527 - bad Masker interface
        #
        # The predicted masks are in the shape of i.e. [N, 1, 28, 28] where N is
        # the number of instances predicted, and they represent the interior
        # of the bounding boxes.
        #
        # Masker projects these predictions on an empty canvas with the full
        # size of the input image using the predicted bounding boxes
        maskTensor = Masker(threshold=0.5).forward_single_image(
            maskTensor, perImagePredictions
        )[:, 0, :, :]

        pixelCounts = []
        for (xmin, ymin, xmax, ymax), instanceMask in zip(bbs, maskTensor):
            pixelCounts.append(instanceMask[ymin:ymax, xmin:xmax].sum().item())

    for predID in range(len(perImagePredictions)):
        xmin, ymin, xmax, ymax = bbs[predID]
        # if we have instance segmentation prediction then we update pixelCount
        pixelCount = 0
        if maskTensor[0] is not None:
            pixelCount = pixelCounts[predID]
            if pixelCount == 0:
                continue

        predInstance = {
            "imgName": idx,
            "predID": predID,
            "labelID": labels[predID],
            "boxArea": boxAreas[predID],
            "pixelCount": pixelCount,
            "confidence": scores[predID],
            "box": (xmin, ymin, xmax, ymax),
            "matchedGt": [],
        }
        perImageInstances.append(predInstance)

    return perImageInstances, maskTensor