Mask-R-CNN生成的mask掩码图像处理

Mask-R-CNN生成的mask掩码图像处理

文章内容

将mask掩码图与原图混合，显示实例分割效果图,省略边框显示

效果如下,实现左图到右图

origin result

数据介绍

原始图像

origin

Mask-R-CNN的mask层输出结果prediction如下

参数解释

boxes:实例分割的边框 dtype=torch.float32
labels对应实例的类别标签 dtype=torch.int64
scores该实例的得分 dtype=torch.float32
masks掩码图图像数组 dtype=torch.float32

上述图像中，Mask-R-CNN网络一共产生了22个可能的mask值

数据处理

数据输入与输出

模型预测代码如下

# 将模型设置为评估模式
# img为输入模型的预测图像
model.eval()
with torch.no_grad():
    prediction = model([img.to(device)])
# 预测结果
prediction=prediction[0]
# 取出img
origin=Image.fromarray(img.mul(255).permute(1, 2, 0).byte().numpy())#原始图像

筛选数据

根据scores的值进行选择mask,此处筛选得分大于0.8的mask

#移动到cpu后转为nunpy格式，便于后续操作
scores=prediction['scores'].cpu().numpy()
score_d=0.8
#得到得分大于score_d掩码数组
mask_np=(scores>score_d)
#读取numpy数组格式的masks，labels
#mul(255)将值归至0-255间
masks=prediction['masks'].mul(255).byte().cpu().numpy()[mask_np]
labels=prediction['labels'].cpu().numpy()[mask_np].astype(np.uint8)

掩码图像转换

定义颜色数组colors

label	color	解释
0	(0, 0, 0)	背景
1	(0 ,255 ,0)	实例类别，如猫
1	(0 ,255, 255)	实例类别，如狗

colors=[(0, 0, 0), (0 ,255 ,0), (0 ,255, 255), (255 ,255, 0), (255, 140, 0), (255, 0, 0),(0, 128, 128)]

遍历masks、labels获得所有的掩码图，并将其转为RGB图像

mask_images=[] #存储所有掩码图
for mask,label in zip(masks,labels):
    #将mask图像复制，输出的mask灰度图像中，值越大颜色越亮，将其与RGB颜色相乘 
    mask_origin=mask.copy().repeat(3,axis=0).transpose(1,2,0) #(1,h,w)->(3,h,w)->(h,w,3)
    #全部值设置为label
    mask[mask>0]=label
    #转为RGB图像，*(mask_o/255)用于调节每个像素点亮度
    #(h,w,1)->(h*w*1,1)->(h*w*1,3)->(h,w,3)
    #(h,w,3)*(h,w,3)
    mask_image = np.reshape(np.array(colors, np.uint8)[np.reshape(mask, [-1])], [origin.size[1], origin.size[0], -1])*(mask_origin/255) 
    mask_images.append(mask_image)

上述步骤单次循环的结果如下，将左图(网络输出的mask灰度图)变成右图(RGB图)，样例中label=1

mask mask_1

掩码图整合

对上述多个实例进行整合到一张图片上

#定义图像大小
result=np.zeros((origin.size[1], origin.size[0],3))
for image in mask_images:
    #融合，每张实例图占的比例一样
    result=result+1/len(mask_images)*image

保存整合的实例图像

cv2.imwrite('masks.png',result)

上述步骤实现的功能效果如下

masks.png

masks

原图与掩码图整合

将掩码图与原图融合

#Image中图像为RGB,opencv中图像为BGR,将通道交换
origin_np=np.array(origin).astype(dtype=np.uint8)[:,:,::-1]
#0.5融合值
result=result*0.5+origin_np*0.5
#保存结果
cv2.imwrite('result.png',result)

result.png

result

整合过程俩者图像会变暗，提升图像整体亮度

# 提高亮度，可以通过加一个常数或者乘以一个倍数来实现
brightness_factor = 1.5  # 调整这个因子来改变亮度
brightened_image = np.clip(result * brightness_factor, 0, 255).astype(np.uint8)
cv2.imwrite('result_up.png',brightened_image)

result_up.png

result_up