从 detectorron2 对象检测分割中提取掩模值，然后用 opencv 绘制掩模并计算该掩模的面积？

Question

我想根据 Detectron2 对象检测分割的输出计算预测掩模的面积 所以当我运行推理时它返回字典

outputs = predictor(im)

pred_mask,pred_boxes,pred_scores

。当我打印

pred_masks

时，值是二进制的

(False,True)

但是当我通过发送预测来调用可视化工具函数时，它会粘贴所有预测的蒙版，但它向我显示（true false）我不明白 这是可视化工具的代码

v = Visualizer(im\[:, :, ::-1\],
metadata=grain_metadata,
scale=0.5,
instance_mode=ColorMode.IMAGE_BW
)
out = v.draw_instance_predictions(outputs\["instances"\].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])

这个函数在一个文件中

(detectron2->utils->visualizer->draw_instance_predictions(predictions)

我在协作中对该文件做了一些更改，例如尝试打印蒙版，但这根本没有影响我对整个文件进行评论，但可视化工具仍在工作，有人可以告诉我如何获取蒙版值，以便我可以借鉴我自己使用 OpenCV。

Answer 1

pred_masks

值中的值是二进制的，因为它包含二进制掩码。

参考Detectron2模型输出格式，可以使用OpenCV的findContours得到每个预测mask的轮廓。

import cv2
import numpy as np

# "outputs" is the inference output in the format described here - https://detectron2.readthedocs.io/tutorials/models.html#model-output-format

# Extract the contour of each predicted mask and save it in a list
contours = []
for pred_mask in outputs[0]['instances'].pred_masks:
    # pred_mask is of type torch.Tensor, and the values are boolean (True, False)
    # Convert it to a 8-bit numpy array, which can then be used to find contours
    mask = pred_mask.numpy().astype('uint8')
    contour, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)

    contours.append(contour[0]) # contour is a tuple (OpenCV 4.5.2), so take the first element which is the array of contour points

然后可以使用

contours

使用 OpenCV 的 drawContours 来绘制对象，如下所示。

# "image" is the original BGR image

image_with_overlaid_predictions = image.copy()

for contour in contours:
    cv2.drawContours(image_with_overlaid_predictions, [contour], -1, (0,255,0), 1)

Answer 2

对 Zepman 的答案进行细微修改：

• 创建一个原始图像大小的空蒙版；
• 迭代输出['instances'].pred_masks，而不是输出[0]；
• 转换前将掩码重新定位到CPU；
• 计算面积；

import cv2
import numpy as np
    
# get dimensions and create an empty mask
height, width, channels = im.shape
binary_mask = np.zeros((height, width, channels), dtype=np.uint8)
    
contours = []
for pred_mask in outputs['instances'].pred_masks:
    mask = pred_mask.cpu().numpy().astype('uint8')
    contour, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    contours.append(contour[0])
    
# draw each contour to the empty mask
for contour in contours:
    cv2.drawContours(binary_mask, [contour], -1, (255,255,255), thickness=cv2.FILLED)

# store the areas in a list
areas = [cv2.contourArea(contour) for contour in contours]

也可以在GIMP或ImageJ等软件中找到导入遮罩的区域。