mmdetection3.0 实例分割结果不同类别分别可视化

蔡不菜和他的uU们

已于 2024-05-17 10:30:40 修改

阅读量686

点赞数 15

分类专栏：机器学习文章标签： deep learning 计算机视觉

于 2024-05-16 14:38:17 首次发布

本文链接：https://blog.csdn.net/csy1021/article/details/138960764

版权

机器学习专栏收录该内容

8 篇文章 0 订阅

订阅专栏

昨天有群友，在群里提到了这个需求，进行实现一下

也就是
一张图片会检测到不同类别的实例，实现的效果就是，将同一类别的输出在一张图片上
例，一张图片检测到 3种实例，输出 3张图片，每一张图上只可视化一种类别

涉及到的文件是，mmdet/visualization/local_visualizer.py

主要逻辑

def add_datasample():
    # ...
    if draw_pred and data_sample is not None:
        pred_img_data = image
        if 'pred_instances' in data_sample:
            pred_instances = data_sample.pred_instances
            pred_instances = pred_instances[
                pred_instances.scores > pred_score_thr]
            pred_img_data = self._draw_instances(image, pred_instances,
                                                 classes, palette)
    # ...
    self.set_image(drawn_img)
    # 写出
    if out_file is not None:
        mmcv.imwrite(drawn_img[..., ::-1], out_file)

def _draw_instances():
    # ...

    # 绘制 检查框 和 标签
    if 'bboxes' in instances and instances.bboxes.sum() > 0:
        self.draw_bboxes()
        for i, (pos, label) in enumerate(zip(positions, labels)):
            self.draw_text()

    # 绘制 实例 mask
    if 'masks' in instances:
        # ...
        self.draw_polygons(polygons, edge_colors='w', alpha=self.alpha)
        self.draw_binary_masks(masks, colors=colors, alphas=self.alpha)

    # ...

修改可以在 _draw_instances() 里面，也可以在 add_datasample() 里面完成

在 add_datasample() 里面修改需要重新拿到 pred_instances，然后进行分析，再绘制，写出
在 _draw_instances() 里面修改进一步传入的pre_instances ,然后返回绘制后的不同类别可视化图像数据列表，在_add_datasample() 写出

这里选择的方式是接近第二种，
根据 label 建立不同的 index 数组，得到每一个类别要绘制的内容，然后调用add_datasample() 进行绘制，不对_draw_instances()

修改 mmdet/visualization/local_visualizer.py

@master_only
def add_datasample(
        self,
        name: str,
        image: np.ndarray,
        data_sample: Optional['DetDataSample'] = None,
        draw_gt: bool = True,
        draw_pred: bool = True,
        show: bool = False,
        wait_time: float = 0,
        # TODO: Supported in mmengine's Viusalizer.
        out_file: Optional[str] = None,
        pred_score_thr: float = 0.3,
        step: int = 0) -> None:

    image = image.clip(0, 255).astype(np.uint8)
    classes = self.dataset_meta.get('classes', None)
    palette = self.dataset_meta.get('palette', None)

    gt_img_data = None
    pred_img_data = None

    if data_sample is not None:
        data_sample = data_sample.cpu()

    if draw_gt and data_sample is not None:
        gt_img_data = image
        if 'gt_instances' in data_sample:
            gt_img_data = self._draw_instances(image,
                                               data_sample.gt_instances,
                                               classes, palette)
        if 'gt_sem_seg' in data_sample:
            gt_img_data = self._draw_sem_seg(gt_img_data,
                                             data_sample.gt_sem_seg,
                                             classes, palette)

        if 'gt_panoptic_seg' in data_sample:
            assert classes is not None, 'class information is ' \
                                        'not provided when ' \
                                        'visualizing panoptic ' \
                                        'segmentation results.'
            gt_img_data = self._draw_panoptic_seg(
                gt_img_data, data_sample.gt_panoptic_seg, classes, palette)
    ########
    draw_image_for_each_class = True
    ########
    if draw_pred and data_sample is not None:
        pred_img_data = image
        if 'pred_instances' in data_sample:
            pred_instances = data_sample.pred_instances
            pred_instances = pred_instances[
                pred_instances.scores > pred_score_thr]
            pred_img_data = self._draw_instances(image, pred_instances,
                                                 classes, palette)
            ################
            each_class_instances_image_list = self._draw_each_class_instances(image, pred_instances,
                                                 classes, palette)
            ################

        if 'pred_sem_seg' in data_sample:
            pred_img_data = self._draw_sem_seg(pred_img_data,
                                               data_sample.pred_sem_seg,
                                               classes, palette)

        if 'pred_panoptic_seg' in data_sample:
            assert classes is not None, 'class information is ' \
                                        'not provided when ' \
                                        'visualizing panoptic ' \
                                        'segmentation results.'
            pred_img_data = self._draw_panoptic_seg(
                pred_img_data, data_sample.pred_panoptic_seg.numpy(),
                classes, palette)

    if gt_img_data is not None and pred_img_data is not None:
        drawn_img = np.concatenate((gt_img_data, pred_img_data), axis=1)
    elif gt_img_data is not None:
        drawn_img = gt_img_data
    elif pred_img_data is not None:
        drawn_img = pred_img_data
    else:
        # Display the original image directly if nothing is drawn.
        drawn_img = image

    # It is convenient for users to obtain the drawn image.
    # For example, the user wants to obtain the drawn image and
    # save it as a video during video inference.
    self.set_image(drawn_img)

    if show:
        self.show(drawn_img, win_name=name, wait_time=wait_time)
    #########
    if draw_image_for_each_class and show:
        for each_class_image in each_class_instances_image_list:
            self.show(each_class_image, win_name=name, wait_time=wait_time)
    #########

    if out_file is not None:
        mmcv.imwrite(drawn_img[..., ::-1], out_file)
        ##############
        if draw_image_for_each_class:

            for i,each_class_image in enumerate(each_class_instances_image_list):
                out_file_ = out_file
                basename = os.path.basename(out_file_)
                out_file_ = os.path.split(out_file_)[0] + '/'+os.path.splitext(basename)[0] + "_" + classes[i] + \
                            os.path.splitext(basename)[-1]
                print(out_file_)
                mmcv.imwrite(each_class_image[..., ::-1],out_file_)
        ##############


    else:
        self.add_image(name, drawn_img, step)

添加、修改的代码一共有四处，用
#############
code
#############
进行了包括，建议不要直接替换，可能存在版本不一致，出现其他问题，对比后，进行添加修改

添加 _draw_each_class_instances()

函数位置放在 _draw_instances() 下面

# only for instance segmentation
def _draw_each_class_instances(self, image: np.ndarray, instances: ['InstanceData'],
                    classes: Optional[List[str]],
                    palette: Optional[List[tuple]]) -> [np.ndarray]:
    each_class_instance_image_list = []
    label_num = 0
    if 'labels' in instances and instances.labels.size()[0] > 0:
        label_num = instances.labels.size()[0]
    if label_num == 0:
        return [image]
    labels = instances.labels.detach().numpy()
    for label in sorted(list(set(labels))):
        index_list = labels == label

        instances_for_label = InstanceData()
        instances_for_label['bboxes'] = instances['bboxes'][torch.tensor(index_list)]
        instances_for_label['labels'] = instances['labels'][torch.tensor(index_list)]
        instances_for_label['masks'] = instances['masks'][torch.tensor(index_list)]
        instances_for_label['scores'] = instances['scores'][torch.tensor(index_list)]
        # instances_for_label['metainfo'] = instances['metainfo']


        each_class_instance_image = self._draw_instances(image,instances_for_label, classes, palette)
        each_class_instance_image_list.append(each_class_instance_image)
    return each_class_instance_image_list