masa/visualization/visualizer.py

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Optional, Tuple, Union

import mmcv
import numpy as np

try:
    import seaborn as sns
except ImportError:
    sns = None
from mmdet.registry import VISUALIZERS
from mmdet.structures import DetDataSample
from mmdet.structures.mask import bitmap_to_polygon
from mmdet.visualization.palette import _get_adaptive_scales
from mmengine.dist import master_only
from mmengine.structures import InstanceData, PixelData
from mmengine.visualization import Visualizer


def random_color(seed):
    """Random a color according to the input seed."""
    if sns is None:
        raise RuntimeError(
            "motmetrics is not installed,\
                 please install it by: pip install seaborn"
        )
    np.random.seed(seed)
    colors = sns.color_palette("tab20")
    color = colors[np.random.choice(range(len(colors)))]
    color = tuple([int(255 * c) for c in color])
    return color


@VISUALIZERS.register_module()
class MasaTrackLocalVisualizer(Visualizer):
    """Tracking Local Visualizer for the MOT, VIS tasks.

    Args:
        name (str): Name of the instance. Defaults to 'visualizer'.
        image (np.ndarray, optional): the origin image to draw. The format
            should be RGB. Defaults to None.
        vis_backends (list, optional): Visual backend config list.
            Defaults to None.
        save_dir (str, optional): Save file dir for all storage backends.
            If it is None, the backend storage will not save any data.
        line_width (int, float): The linewidth of lines.
            Defaults to 3.
        alpha (int, float): The transparency of bboxes or mask.
                Defaults to 0.8.
    """

    def __init__(
        self,
        name: str = "visualizer",
        image: Optional[np.ndarray] = None,
        vis_backends: Optional[Dict] = None,
        save_dir: Optional[str] = None,
        line_width: Union[int, float] = 5,
        alpha: float = 0.8,
        show_polygon: bool = False,
        texts=None,
    ) -> None:
        super().__init__(name, image, vis_backends, save_dir)
        self.line_width = line_width
        self.alpha = alpha
        self.show_polygon = show_polygon
        # Set default value. When calling
        # `TrackLocalVisualizer().dataset_meta=xxx`,
        # it will override the default value.
        self.dataset_meta = {}
        if texts is not None:
            if isinstance(texts, str):
                if not texts.endswith("."):
                    original_caption = texts + " . "
                original_caption = original_caption.split(" . ")
                class_names = list(filter(lambda x: len(x) > 0, original_caption))
            else:
                class_names = list(texts)
            self.label_names = class_names
        else:
            self.label_names = None

    def _draw_instances(
        self, image: np.ndarray, instances: ["InstanceData"]
    ) -> np.ndarray:
        """Draw instances of GT or prediction.

        Args:
            image (np.ndarray): The image to draw.
            instances (:obj:`InstanceData`): Data structure for
                instance-level annotations or predictions.
        Returns:
            np.ndarray: the drawn image which channel is RGB.
        """
        self.set_image(image)
        classes = self.dataset_meta.get("classes", None)

        # get colors and texts
        # for the MOT and VIS tasks
        colors = [random_color(_id) for _id in instances.instances_id]

        # draw bboxes and texts
        if "bboxes" in instances:
            # draw bboxes
            bboxes = instances.bboxes.clone()
            labels = instances.labels.clone()

            self.draw_bboxes(
                bboxes,
                edge_colors=colors,
                alpha=self.alpha,
                line_widths=self.line_width,
            )
            # draw texts
            positions = bboxes[:, :2] - self.line_width
            areas = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0])
            scales = _get_adaptive_scales(areas.cpu().numpy())

            for i, (pos, label) in enumerate(zip(positions, labels)):
                if self.label_names is not None:
                    label_text = self.label_names[label]
                else:
                    label_text = (
                        classes[label] if classes is not None else f"class {label}"
                    )

                if "instances_id" in instances:
                    label_text += f" | {instances.instances_id[i]}"

                if "scores" in instances:
                    score = round(float(instances.scores[i]) * 100, 1)
                    label_text += f": {score}"

                self.draw_texts(
                    label_text,
                    pos,
                    colors="black",
                    font_sizes=int(13 * scales[i]),
                    bboxes=[
                        {
                            "facecolor": [c / 255 for c in colors[i]],
                            "alpha": 0.8,
                            "pad": 0.7,
                            "edgecolor": "none",
                        }
                    ],
                )

        # draw masks
        if "masks" in instances:
            masks = instances.masks
            polygons = []
            for i, mask in enumerate(masks):
                contours, _ = bitmap_to_polygon(mask)
                polygons.extend(contours)
            if self.show_polygon:
                self.draw_polygons(polygons, edge_colors="w", alpha=self.alpha)
            self.draw_binary_masks(masks, colors=colors, alphas=self.alpha)

        return self.get_image()

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

        - If GT and prediction are plotted at the same time, they are
        displayed in a stitched image where the left image is the
        ground truth and the right image is the prediction.
        - If ``show`` is True, all storage backends are ignored, and
        the images will be displayed in a local window.
        - If ``out_file`` is specified, the drawn image will be
        saved to ``out_file``. t is usually used when the display
        is not available.
        Args:
            name (str): The image identifier.
            image (np.ndarray): The image to draw.
            data_sample (OptTrackSampleList): A data
                sample that contain annotations and predictions.
                Defaults to None.
            draw_gt (bool): Whether to draw GT TrackDataSample.
                Default to True.
            draw_pred (bool): Whether to draw Prediction TrackDataSample.
                Defaults to True.
            show (bool): Whether to display the drawn image. Default to False.
            wait_time (int): The interval of show (s). Defaults to 0.
            out_file (str): Path to output file. Defaults to None.
            pred_score_thr (float): The threshold to visualize the bboxes
                and masks. Defaults to 0.3.
            step (int): Global step value to record. Defaults to 0.
        """
        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:
            assert "gt_instances" in data_sample
            gt_img_data = self._draw_instances(image, data_sample.gt_instances)

        if draw_pred and data_sample is not None:
            assert "pred_track_instances" in data_sample
            pred_instances = data_sample.pred_track_instances
            if "scores" in pred_instances:
                pred_instances = pred_instances[
                    pred_instances.scores > pred_score_thr
                ].cpu()
            pred_img_data = self._draw_instances(image, pred_instances)

        if fps is not None:
            self.draw_texts(
                f"FPS: {fps: .1f}",
                np.array([10, 10]),
                colors="black",
                font_sizes=15,
                bboxes=[
                    {"facecolor": "w", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"}
                ],
            )

        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
        else:
            drawn_img = pred_img_data

        if show:
            self.show(drawn_img, win_name=name, wait_time=wait_time)

        if out_file is not None:
            mmcv.imwrite(drawn_img[..., ::-1], out_file)
        else:
            self.add_image(name, drawn_img, step)