python - 用新实例替换类的实例

Question

我正在研究 detectron2 对象检测。我在过滤检测到的对象时遇到问题。

这里是 detectron2 的预测输出:

Instances(num_instances=9, image_height=547, image_width=820, fields=[pred_boxes: Boxes(tensor([[3.1173e+01, 3.8368e+01, 5.3751e+02, 5.4078e+02],
    [5.9945e+02, 2.6412e+02, 6.8196e+02, 5.1333e+02],
    [4.4486e+02, 1.7210e+02, 4.9981e+02, 2.5596e+02],
    [1.1566e-01, 2.3533e+02, 8.5483e+01, 3.6838e+02],
    [3.0897e+02, 2.4964e+02, 3.5739e+02, 4.8948e+02],
    [7.6962e-03, 2.3240e+02, 8.5447e+01, 3.7128e+02],
    [2.7454e+02, 2.6212e+02, 3.3122e+02, 4.5928e+02],
    [6.4399e+02, 3.0057e+02, 6.6374e+02, 3.8033e+02],
    [3.1025e+02, 2.5372e+02, 3.3572e+02, 3.5059e+02]])), scores: tensor([0.9998, 0.9994, 0.9941, 0.8815, 0.8447, 0.3559, 0.1484, 0.1304, 0.0928]), pred_classes: tensor([ 0,  0, 67,  2, 27,  7, 27, 27, 27]), pred_masks: tensor([[[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]],

    [[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]],

    [[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]],

    ...,

    [[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]],

    [[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]],

    [[False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False]]])])

我进行了过滤并创建了一个包含预测对象类、分数和框的新列表 (dict)。我想在图像上绘制和可视化它:

过滤代码:

idxofClass = [i for i, x in enumerate(list(outputs['instances'].pred_classes)) if (x == 0)]
outputs_new = [{'pred_classes': o.pred_classes[idxofClass], 'scores':o.scores[idxofClass], 'pred_boxes':o.pred_boxes[idxofClass] }]

现在，我可以得到过滤后的值如下:

[{'pred_classes': tensor([ 0,  0, 67]), 'scores': tensor([0.9998, 0.9994, 0.9941]), 'pred_boxes': Boxes(tensor([[ 31.1728,  38.3685, 537.5092, 540.7788],
    [599.4498, 264.1228, 681.9622, 513.3326],
    [444.8603, 172.1017, 499.8055, 255.9632]]))}]

将此值传递给 Visualizer 时，出现以下错误:

Traceback (most recent call last):
  File "apimodel.py", line 96, in <module>
    out = v.draw_instance_predictions(outputs_new)
  File "/root/anaconda3/envs/ml-engine/lib/python3.8/site-packages/detectron2/utils/visualizer.py", line 366, in draw_instance_predictions
    boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None
AttributeError: 'list' object has no attribute 'has'

原始输出的数据类型是类实例:

o = outputs["instances"]
print("data type:", type(o))
<class 'detectron2.structures.instances.Instances'>

新创建的过滤输出的输出是一个list(dict):

<class 'list'>

我的目标是根据过滤分数绘制边界框。我一直在尝试替换输出的原始值，但没有成功。请为此提供帮助。

Answer 1

经过两天的搜索，我找到了实现目标的方法。我正在写答案，作为 detectron2 类，这样，如果任何人寻找类似的方法都会受益。

类的过滤索引:

idxofClass = [i for i, x in enumerate(list(outputs['instances'].pred_classes)) if x == 0]

创建新类、框、分数和 mask :

classes = o.pred_classes[idxofClass]
scores = o.scores[idxofClass]
boxes = o.pred_boxes[idxofClass]
masks = o.pred_masks[idxofClass]

定义新实例并将新值设置为新实例。注意:detectron2模块提供了这个方法set。

obj = detectron2.structures.Instances(image_size=(480, 640))

obj.set('pred_classes', classes)
obj.set('scores', scores)
obj.set('pred_boxes', boxes)
obj.set('pred_masks', masks)

现在您可以使用这个新实例 obj 进行其他处理和可视化:

out = v.draw_instance_predictions(obj.to("cpu"))