在微调Faster RCNN目标检测模型之后，如何可视化边界框预测？

Question

我使用自定义数据集fine-tune了pytorch torchvision模型=torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)。

我按照这份指南只训练了Faster RCNN，而不是Mask RCNN。

训练期间没有出现任何错误，模型返回了一个包含预测框、标签和得分的字典。

在我所遵循的指南中，他们展示了如何可视化模型预测的掩膜。是否有类似的方法来可视化边界框？我在解决这个问题时遇到了很多麻烦。

谢谢。

- adam l

Answer 1

你可以使用FiftyOne轻松创建数据集并添加预测结果。然后，您可以在交互式应用程序中可视化它们。本教程遵循与您之前所做的类似的工作流程：https://voxel51.com/docs/fiftyone/tutorials/evaluate_detections.html。这比单独可视化每个图像要好得多。

- Eric Hofesmann

---

我们发现FiftyOne是一个改变游戏规则的工具。它消除了在https://dev59.com/WLnoa4cB1Zd3GeqPOFDH#60275224中描述的可视化编码需求。强烈推荐使用。 - mherzog

Answer 2

FasterRCNN的预测结果形式为：

>>> predictions = model([input_img_tensor])
[{'boxes': tensor([[419.6865, 170.0683, 536.0842, 493.7452],
          [159.0727, 180.3606, 298.8194, 434.4604],
          [439.7836, 222.6208, 452.0138, 271.8359],
          [444.3562, 224.4628, 456.1511, 265.5336],
          [437.7808, 226.5965, 446.2904, 271.2691]], grad_fn=<StackBackward>),
  'labels': tensor([ 1,  1, 32, 32, 32]),
  'scores': tensor([0.9997, 0.9996, 0.5827, 0.2102, 0.0943], grad_fn=<IndexBackward>)}]

预测的框是以[x1，y1，x2，y2]格式表示的，其值介于0和H以及0和W之间。

您可以使用OpenCV的rectangle函数在图像上叠加边界框。

import cv2
img = cv2.imread('input_iamge.png', cv2.COLOR_BGR2RGB)

for i in range(len(predictions[0]['boxes'])):
    x1, x2, x3, x4 = map(int, predictions[0]['boxes'][i].tolist())
    print(x1, x2, x3, x4)
    image = cv2.rectangle(img, (x1, x2), (x3, x4), (255, 0, 0), 1)

cv2_imshow('img', image)

- kHarshit

---

NameError：名称'cv2_imshow'未定义，对于cv2版本4.6.0，请详细说明如何修复较新的python/cv2版本中的此错误。谢谢。 - 7shoe

@7shoe 使用 cv2.imshow() 代替。cv2_imshow() 只在 colab 上可用。 - kHarshit