从 Mask RCNN 结果中选择一个特定的掩码,r。过滤字典以仅保留 r['class_ids']==1 处的键

Selecting a specific mask from Mask RCNN results, r. Filtering a dictionary to only keep key where r['class_ids']==1

这是最接近我现在遇到的问题的问题:

掩码 RCNN returns 字典,r,看起来像:

{'rois': array([[ 17,   0, 371, 300],
        [323, 107, 374, 153]]),
 'class_ids': array([ 1, 28]),
 'scores': array([0.9999876 , 0.77787864], dtype=float32),
 'masks': array([[[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]]])}

我只想为那个人保留口罩class (class_ids == 1)

我试图通过

过滤字典
filtered = {k: v for k, v in r.items() if (r['class_ids'] == 1)}

这导致了以下 ValueError

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

当我再次尝试时

filtered = {k: v for k, v in r.items() if (r['class_ids'] == 1).any()}

这不会过滤掉任何东西。 a.all() returns 和空字典。我很困惑该怎么做。

如有任何帮助,我们将不胜感激!谢谢!

根据我的理解,r['masks'] 数组包含模型中每个可能 class 的数组。使用 COCO 作为基础模型。 person class 的索引为 1。因此您可以使用

获取 person mask
person_mask = r['masks'][:][1]

但是,在我的用例中,它并没有多大帮助。因为我想将人物面具放在白色背景上。使用 skimage.io.imread 读取的原始图像具有形状 (x, y, z) 并且 person_mask 具有形状 (y, z) 因此最好不要过滤 r['masks'] ,预先具有形状(x,y,z)。 r['masks'] 形状为 (x, y, z) 其中 z 是在图像中找到的面具数量 - 例如它可以是人的面具(索引 1)和领带的面具(索引 28)。通常,假设第一个蒙版是人的蒙版是安全的,这就是为什么我们可以使用

将原始图像转换为仅保留人蒙版的区域
temp[:,:,j] = temp[:,:,j] * r['masks'][:,:,0]

完整代码如下

temp = skimage.io.imread('05a.jpg')
for j in range(temp.shape[2]):
    temp[:,:,j] = temp[:,:,j] * r['masks'][:,:,0]
    background = np.zeros_like(temp)
    background[:,:,:]  = (255, 255, 255)
    rows,cols,channels = temp.shape
    roi = background[0:rows, 0:cols]

    temp = cv.cvtColor(temp,cv.COLOR_RGB2BGR)
    img2gray = cv.cvtColor(temp,cv.COLOR_BGR2GRAY)
    ret, mask = cv.threshold(img2gray, 0, 255, cv.THRESH_BINARY)
    mask_inv = cv.bitwise_not(mask)

    img1_bg = cv.bitwise_and(roi,roi,mask = mask_inv)

    img2_fg = cv.bitwise_and(temp,temp,mask = mask)

    dst = cv.add(img1_bg,img2_fg)
    background[0:rows, 0:cols ] = dst
    cv.imwrite('res.jpg', background)