问题背景

训练了Detectron2 中的实例分割网络，是二分类（背景和目标），实验主要是做语义分割的，所以想输出二值化（0和255）的mask掩码，来计算指标。

解决方法

步骤一

在 demo/predictor.py 中作修改：

添加的代码如下：

# 制作mask掩码
pred = (predictions["instances"]._fields)["pred_masks"].cpu().numpy() # bool 类型
binary_mask = np.zeros((pred.shape[1], pred.shape[2]))
for i in range(pred.shape[0]):
    binary_mask = binary_mask + pred[i,:,:]
binary_mask[binary_mask > 0] = 255

步骤二

将返回的二值图像保存，需要修改 tool/demo.py ：

附录

完整的 run_on_image() 函数如下，返回值中添加了二值图像：

    def run_on_image(self, image):
        """
        Args:
            image (np.ndarray): an image of shape (H, W, C) (in BGR order).
                This is the format used by OpenCV.

        Returns:
            predictions (dict): the output of the model.
            vis_output (VisImage): the visualized image output.
        """
        vis_output = None
        predictions = self.predictor(image)

        # 制作mask掩码
        pred = (predictions["instances"]._fields)["pred_masks"].cpu().numpy()
        binary_mask = np.zeros((pred.shape[1], pred.shape[2]))
        for i in range(pred.shape[0]):
            binary_mask = binary_mask + pred[i,:,:]
        binary_mask[binary_mask > 0] = 255
        # print(binary_mask.shape)

        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        image = image[:, :, ::-1]
        visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)
        if "panoptic_seg" in predictions:
            panoptic_seg, segments_info = predictions["panoptic_seg"]
            vis_output = visualizer.draw_panoptic_seg_predictions(
                panoptic_seg.to(self.cpu_device), segments_info
            )
        else:
            if "sem_seg" in predictions:
                vis_output = visualizer.draw_sem_seg(
                    predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)
                )
            if "instances" in predictions:
                instances = predictions["instances"].to(self.cpu_device)
                vis_output = visualizer.draw_instance_predictions(predictions=instances)

        return predictions, vis_output, binary_mask