[CodeCamp2023-154] Add semantic label to the segmentation visualization results (open-mmlab#3229)

Thanks for your contribution and we appreciate it a lot. The following
instructions would make your pull request more healthy and more easily
get feedback. If you do not understand some items, don't worry, just
make the pull request and seek help from maintainers.

## Motivation

[Add semantic label to the segmentation visualization results
分割可视化结果中加上语义信息
open-mmlab#154](open-mmlab/OpenMMLabCamp#154)

corresponding issue: [跑出来结果之后怎么在结果图片上获取各个语意部分的区域信息？
open-mmlab#2578](open-mmlab#2578)

## Modification

1. mmseg/apis/inference.py, add withLabels in visualizer.add_datasample
call, to indicate whether add semantic label
2. mmseg/visualization/local_visualizer.py, add semantic labels by
opencv; modify the demo comment description
3. mmseg/utils/__init__.py, add bdd100k datasets to test
local_visualizer.py

**Current visualize result**
<img width="637" alt="image"
src="https://github.com/open-mmlab/mmsegmentation/assets/35064479/6ef6ce02-1d82-46f8-bde9-a1d69ff62df8">


**Add semantic label**
<img width="637" alt="image"
src="https://github.com/open-mmlab/mmsegmentation/assets/35064479/00716679-b43a-4794-8499-9bfecdb4b78b">

## Test results
**tests/test_visualization/test_local_visualizer.py** test
results：（MMSegmentation/tests/data/pseudo_cityscapes_dataset/leftImg8bit/val/frankfurt/frankfurt_000000_000294_leftImg8bit.png）
<img width="643" alt="image"
src="https://github.com/open-mmlab/mmsegmentation/assets/35064479/6792b7d2-2512-4ea9-8500-1a7ed2d5e0dc">

**demo/inference_demo.ipynb** test results：
<img width="966" alt="image"
src="https://github.com/open-mmlab/mmsegmentation/assets/35064479/dfc0147e-fb1a-490a-b6ff-a8b209352d9b">

-----
## Drawbacks
config opencv thickness according to image size
<img width="496" alt="image"
src="https://github.com/open-mmlab/mmsegmentation/assets/35064479/0a54d72c-62b1-422c-89ae-69dc753fe0fc">

I have no idea of dealing with label overlapping for the time being

Author: CastleDream

mmseg/apis/inference.py

@@ -158,6 +158,7 @@ def show_result_pyplot(model: BaseSegmentor,
                        draw_pred: bool = True,
                        wait_time: float = 0,
                        show: bool = True,
+                       withLabels: Optional[bool] = True,
                        save_dir=None,
                        out_file=None):
     """Visualize the segmentation results on the image.
@@ -177,10 +178,14 @@ def show_result_pyplot(model: BaseSegmentor,
             that means "forever". Defaults to 0.
         show (bool): Whether to display the drawn image.
             Default to True.
+        withLabels(bool, optional): Add semantic labels in visualization
+            result, Default to True.
         save_dir (str, optional): Save file dir for all storage backends.
             If it is None, the backend storage will not save any data.
         out_file (str, optional): Path to output file. Default to None.
 
+
+
     Returns:
         np.ndarray: the drawn image which channel is RGB.
     """
@@ -208,7 +213,8 @@ def show_result_pyplot(model: BaseSegmentor,
         draw_pred=draw_pred,
         wait_time=wait_time,
         out_file=out_file,
-        show=show)
+        show=show,
+        withLabels=withLabels)
     vis_img = visualizer.get_image()
 
     return vis_img

mmseg/utils/__init__.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 # yapf: disable
-from .class_names import (ade_classes, ade_palette, cityscapes_classes,
+from .class_names import (ade_classes, ade_palette, bdd100k_classes,
+                          bdd100k_palette, cityscapes_classes,
                           cityscapes_palette, cocostuff_classes,
                           cocostuff_palette, dataset_aliases, get_classes,
                           get_palette, isaid_classes, isaid_palette,
@@ -27,5 +28,6 @@
     'cityscapes_palette', 'ade_palette', 'voc_palette', 'cocostuff_palette',
     'loveda_palette', 'potsdam_palette', 'vaihingen_palette', 'isaid_palette',
     'stare_palette', 'dataset_aliases', 'get_classes', 'get_palette',
-    'datafrombytes', 'synapse_palette', 'synapse_classes'
+    'datafrombytes', 'synapse_palette', 'synapse_classes', 'bdd100k_classes',
+    'bdd100k_palette'
 ]

mmseg/visualization/local_visualizer.py

@@ -1,8 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from typing import Dict, List, Optional
 
+import cv2
 import mmcv
 import numpy as np
+import torch
 from mmengine.dist import master_only
 from mmengine.structures import PixelData
 from mmengine.visualization import Visualizer
@@ -42,8 +44,8 @@ class SegLocalVisualizer(Visualizer):
         >>> import numpy as np
         >>> import torch
         >>> from mmengine.structures import PixelData
-        >>> from mmseg.data import SegDataSample
-        >>> from mmseg.engine.visualization import SegLocalVisualizer
+        >>> from mmseg.structures import SegDataSample
+        >>> from mmseg.visualization import SegLocalVisualizer
 
         >>> seg_local_visualizer = SegLocalVisualizer()
         >>> image = np.random.randint(0, 256,
@@ -60,7 +62,7 @@ class SegLocalVisualizer(Visualizer):
         >>> seg_local_visualizer.add_datasample(
         ...                        'visualizer_example', image,
         ...                         gt_seg_data_sample, show=True)
-    """ # noqa
+    """  # noqa
 
     def __init__(self,
                  name: str = 'visualizer',
@@ -76,9 +78,32 @@ def __init__(self,
         self.alpha: float = alpha
         self.set_dataset_meta(palette, classes, dataset_name)
 
-    def _draw_sem_seg(self, image: np.ndarray, sem_seg: PixelData,
+    def _get_center_loc(self, mask: np.ndarray) -> np.ndarray:
+        """Get semantic seg center coordinate.
+
+        Args:
+            mask: np.ndarray: get from sem_seg
+        """
+        loc = np.argwhere(mask == 1)
+
+        loc_sort = np.array(
+            sorted(loc.tolist(), key=lambda row: (row[0], row[1])))
+        y_list = loc_sort[:, 0]
+        unique, indices, counts = np.unique(
+            y_list, return_index=True, return_counts=True)
+        y_loc = unique[counts.argmax()]
+        y_most_freq_loc = loc[loc_sort[:, 0] == y_loc]
+        center_num = len(y_most_freq_loc) // 2
+        x = y_most_freq_loc[center_num][1]
+        y = y_most_freq_loc[center_num][0]
+        return np.array([x, y])
+
+    def _draw_sem_seg(self,
+                      image: np.ndarray,
+                      sem_seg: PixelData,
                       classes: Optional[List],
-                      palette: Optional[List]) -> np.ndarray:
+                      palette: Optional[List],
+                      withLabels: Optional[bool] = True) -> np.ndarray:
         """Draw semantic seg of GT or prediction.
 
         Args:
@@ -94,6 +119,8 @@ def _draw_sem_seg(self, image: np.ndarray, sem_seg: PixelData,
             palette (list, optional): Input palette for result rendering, which
                 is a list of color palette responding to the classes.
                 Defaults to None.
+            withLabels(bool, optional): Add semantic labels in visualization
+                result, Default to True.
 
         Returns:
             np.ndarray: the drawn image which channel is RGB.
@@ -112,6 +139,43 @@ def _draw_sem_seg(self, image: np.ndarray, sem_seg: PixelData,
         for label, color in zip(labels, colors):
             mask[sem_seg[0] == label, :] = color
 
+        if withLabels:
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            # (0,1] to change the size of the text relative to the image
+            scale = 0.05
+            fontScale = min(image.shape[0], image.shape[1]) / (25 / scale)
+            fontColor = (255, 255, 255)
+            if image.shape[0] < 300 or image.shape[1] < 300:
+                thickness = 1
+                rectangleThickness = 1
+            else:
+                thickness = 2
+                rectangleThickness = 2
+            lineType = 2
+
+            if isinstance(sem_seg[0], torch.Tensor):
+                masks = sem_seg[0].numpy() == labels[:, None, None]
+            else:
+                masks = sem_seg[0] == labels[:, None, None]
+            masks = masks.astype(np.uint8)
+            for mask_num in range(len(labels)):
+                classes_id = labels[mask_num]
+                classes_color = colors[mask_num]
+                loc = self._get_center_loc(masks[mask_num])
+                text = classes[classes_id]
+                (label_width, label_height), baseline = cv2.getTextSize(
+                    text, font, fontScale, thickness)
+                mask = cv2.rectangle(mask, loc,
+                                     (loc[0] + label_width + baseline,
+                                      loc[1] + label_height + baseline),
+                                     classes_color, -1)
+                mask = cv2.rectangle(mask, loc,
+                                     (loc[0] + label_width + baseline,
+                                      loc[1] + label_height + baseline),
+                                     (0, 0, 0), rectangleThickness)
+                mask = cv2.putText(mask, text, (loc[0], loc[1] + label_height),
+                                   font, fontScale, fontColor, thickness,
+                                   lineType)
         color_seg = (image * (1 - self.alpha) + mask * self.alpha).astype(
             np.uint8)
         self.set_image(color_seg)
@@ -137,7 +201,7 @@ def set_dataset_meta(self,
                 visulizer will use the meta information of the dataset i.e.
                 classes and palette, but the `classes` and `palette` have
                 higher priority. Defaults to None.
-        """ # noqa
+        """  # noqa
         # Set default value. When calling
         # `SegLocalVisualizer().dataset_meta=xxx`,
         # it will override the default value.
@@ -161,7 +225,8 @@ def add_datasample(
             wait_time: float = 0,
             # TODO: Supported in mmengine's Viusalizer.
             out_file: Optional[str] = None,
-            step: int = 0) -> None:
+            step: int = 0,
+            withLabels: Optional[bool] = True) -> None:
         """Draw datasample and save to all backends.
 
         - If GT and prediction are plotted at the same time, they are
@@ -187,6 +252,8 @@ def add_datasample(
             wait_time (float): The interval of show (s). Defaults to 0.
             out_file (str): Path to output file. Defaults to None.
             step (int): Global step value to record. Defaults to 0.
+            withLabels(bool, optional): Add semantic labels in visualization
+                result, Defaults to True.
         """
         classes = self.dataset_meta.get('classes', None)
         palette = self.dataset_meta.get('palette', None)
@@ -202,7 +269,7 @@ def add_datasample(
                                         'segmentation results.'
             gt_img_data = self._draw_sem_seg(gt_img_data,
                                              data_sample.gt_sem_seg, classes,
-                                             palette)
+                                             palette, withLabels)
 
         if (draw_pred and data_sample is not None
                 and 'pred_sem_seg' in data_sample):
@@ -213,7 +280,7 @@ def add_datasample(
                                         'segmentation results.'
             pred_img_data = self._draw_sem_seg(pred_img_data,
                                                data_sample.pred_sem_seg,
-                                               classes, palette)
+                                               classes, palette, withLabels)
 
         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)