Spaces:
Running on Zero
Running on Zero
File size: 8,688 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | """Bounding box visualizer."""
from __future__ import annotations
import os
from dataclasses import dataclass
from mapdet3d.common.typing import (
ArgsType,
ArrayLike,
ArrayLikeFloat,
ArrayLikeInt,
NDArrayUI8,
)
from mapdet3d.vis.base import Visualizer
from mapdet3d.vis.util import generate_color_map
from .canvas import PillowCanvasBackend
from .util import preprocess_boxes, preprocess_image
from .viewer import MatplotlibImageViewer
@dataclass
class DetectionBox2D:
"""Dataclass storing box informations."""
corners: tuple[float, float, float, float]
label: str
color: tuple[int, int, int]
@dataclass
class DataSample:
"""Dataclass storing a data sample that can be visualized."""
image: NDArrayUI8
image_name: str
boxes: list[DetectionBox2D]
class BoundingBoxVisualizer(Visualizer):
"""Bounding box visualizer class."""
def __init__(
self,
*args: ArgsType,
n_colors: int = 50,
cat_mapping: dict[str, int] | None = None,
file_type: str = "png",
width: int = 2,
canvas: PillowCanvasBackend = PillowCanvasBackend(),
viewer: MatplotlibImageViewer = MatplotlibImageViewer(),
**kwargs: ArgsType,
) -> None:
"""Creates a new Visualizer for Image and Bounding Boxes.
Args:
n_colors (int): How many colors should be used for the internal
color map
cat_mapping (dict[str, int]): Mapping from class names to class
ids. Defaults to None.
file_type (str): Desired file type. Defaults to "png".
width (int): Width of the bounding box lines. Defaults to 2.
canvas (CanvasBackend): Backend that is used to draw on images.
viewer (ImageViewerBackend): Backend that is used show images.
"""
super().__init__(*args, **kwargs)
self._samples: list[DataSample] = []
self.color_palette = generate_color_map(n_colors)
self.class_id_mapping = (
{v: k for k, v in cat_mapping.items()}
if cat_mapping is not None
else {}
)
self.file_type = file_type
self.width = width
self.canvas = canvas
self.viewer = viewer
def __repr__(self) -> str:
"""Return string representation of the visualizer."""
return "BoundingBoxVisualizer"
def reset(self) -> None:
"""Reset visualizer."""
self._samples.clear()
def process( # pylint: disable=arguments-differ
self,
cur_iter: int,
images: list[ArrayLike] | None = None,
views: list[list[dict[str, ArrayLike]]] | None = None,
image_names: list[str] | None = None,
boxes: list[ArrayLikeFloat] | None = None,
scores: None | list[ArrayLikeFloat] = None,
class_ids: None | list[ArrayLikeInt] = None,
track_ids: None | list[ArrayLikeInt] = None,
categories: None | list[list[str]] = None,
) -> None:
"""Processes a batch of data.
Args:
cur_iter (int): Current iteration.
images (list[ArrayLike]): Images to show.
image_names (list[str]): Image names.
boxes (list[ArrayLikeFloat]): List of predicted bounding boxes with
shape [N, (x1, y1, x2, y2)], where N is the number of boxes.
scores (None | list[ArrayLikeFloat], optional): List of predicted
box scores each of shape [N]. Defaults to None.
class_ids (None | list[ArrayLikeInt], optional): List of predicted
class ids each of shape [N]. Defaults to None.
track_ids (None | list[ArrayLikeInt], optional): List of predicted
track ids each of shape [N]. Defaults to None.
categories (None | list[list[str]], optional): List of categories
for each image. Instead of class ids, the categories will be
used to label the boxes. Defaults to None.
"""
if self._run_on_batch(cur_iter):
if views is not None:
images = [v[0]["img"] for v in views]
image_names = [n[0] for n in image_names]
boxes = [b[0] for b in boxes]
if categories is not None:
categories = [c[0] for c in categories]
if track_ids is not None:
track_ids = [t[0] for t in track_ids]
for idx, image in enumerate(images):
self.process_single_image(
image,
image_names[idx],
boxes[idx],
None if scores is None else scores[idx],
None if class_ids is None else class_ids[idx],
None if track_ids is None else track_ids[idx],
None if categories is None else categories[idx],
)
def process_single_image(
self,
image: ArrayLike,
image_name: str,
boxes: ArrayLikeFloat,
scores: None | ArrayLikeFloat = None,
class_ids: None | ArrayLikeInt = None,
track_ids: None | ArrayLikeInt = None,
categories: None | list[str] = None,
) -> None:
"""Processes a single image entry.
Args:
image (ArrayLike): Image to show.
image_name (str): Image name.
boxes (ArrayLikeFloat): Predicted bounding boxes with shape
[N, (x1,y1,x2,y2)], where N is the number of boxes.
scores (None | ArrayLikeFloat, optional): Predicted box scores of
shape [N]. Defaults to None.
class_ids (None | ArrayLikeInt, optional): Predicted class ids of
shape [N]. Defaults to None.
track_ids (None | ArrayLikeInt, optional): Predicted track ids of
shape [N]. Defaults to None.
categories (None | list[str], optional): List of categories for
each box. Instead of class ids, the categories will be used to
label the boxes. Defaults to None.
"""
img_normalized = preprocess_image(image, mode=self.image_mode)
data_sample = DataSample(img_normalized, image_name, [])
for corners, label, color in zip(
*preprocess_boxes(
boxes,
scores,
class_ids,
track_ids,
self.color_palette,
self.class_id_mapping,
categories=categories,
)
):
data_sample.boxes.append(
DetectionBox2D(
corners=(corners[0], corners[1], corners[2], corners[3]),
label=label,
color=color,
)
)
self._samples.append(data_sample)
def show(self, cur_iter: int, blocking: bool = True) -> None:
"""Shows the processed images in a interactive window.
Args:
cur_iter (int): Current iteration.
blocking (bool): If the visualizer should be blocking i.e. wait for
human input for each image. Defaults to True.
"""
if self._run_on_batch(cur_iter):
image_data = [self._draw_image(d) for d in self._samples]
self.viewer.show_images(image_data, blocking=blocking)
def _draw_image(self, sample: DataSample) -> NDArrayUI8:
"""Visualizes the datasample and returns is as numpy image.
Args:
sample (DataSample): The data sample to visualize.
Returns:
NDArrayUI8: A image with the visualized data sample.
"""
self.canvas.create_canvas(sample.image)
for box in sample.boxes:
self.canvas.draw_box(box.corners, box.color, width=self.width)
self.canvas.draw_text(box.corners[:2], box.label, box.color)
return self.canvas.as_numpy_image()
def save_to_disk(self, cur_iter: int, output_folder: str) -> None:
"""Saves the visualization to disk.
Writes all processes samples to the output folder naming each image
<sample.image_name>.<filetype>.
Args:
cur_iter (int): Current iteration.
output_folder (str): Folder where the output should be written.
"""
if self._run_on_batch(cur_iter):
for sample in self._samples:
image_name = f"{sample.image_name}.{self.file_type}"
_ = self._draw_image(sample)
self.canvas.save_to_disk(
os.path.join(output_folder, image_name)
)
|