xdecoder/Instruct-X-Decoder
163
1# Copyright (c) Facebook, Inc. and its affiliates.2import colorsys3import logging4import math5import numpy as np6from enum import Enum, unique7import cv28import matplotlib as mpl9import matplotlib.colors as mplc10import matplotlib.figure as mplfigure11import pycocotools.mask as mask_util12import torch13from matplotlib.backends.backend_agg import FigureCanvasAgg14from PIL import Image15 16from detectron2.data import MetadataCatalog17from detectron2.structures import BitMasks, Boxes, BoxMode, Keypoints, PolygonMasks, RotatedBoxes18from detectron2.utils.file_io import PathManager19 20from detectron2.utils.colormap import random_color21 22logger = logging.getLogger(__name__)23__all__ = ["ColorMode", "VisImage", "Visualizer"]24 25 26_SMALL_OBJECT_AREA_THRESH = 100027_LARGE_MASK_AREA_THRESH = 12000028_OFF_WHITE = (1.0, 1.0, 240.0 / 255)29_BLACK = (0, 0, 0)30_RED = (1.0, 0, 0)31 32_KEYPOINT_THRESHOLD = 0.0533 34 35@unique36class ColorMode(Enum):37 """38 Enum of different color modes to use for instance visualizations.39 """40 41 IMAGE = 042 """43 Picks a random color for every instance and overlay segmentations with low opacity.44 """45 SEGMENTATION = 146 """47 Let instances of the same category have similar colors48 (from metadata.thing_colors), and overlay them with49 high opacity. This provides more attention on the quality of segmentation.50 """51 IMAGE_BW = 252 """53 Same as IMAGE, but convert all areas without masks to gray-scale.54 Only available for drawing per-instance mask predictions.55 """56 57 58class GenericMask:59 """60 Attribute:61 polygons (list[ndarray]): list[ndarray]: polygons for this mask.62 Each ndarray has format [x, y, x, y, ...]63 mask (ndarray): a binary mask64 """65 66 def __init__(self, mask_or_polygons, height, width):67 self._mask = self._polygons = self._has_holes = None68 self.height = height69 self.width = width70 71 m = mask_or_polygons72 if isinstance(m, dict):73 # RLEs74 assert "counts" in m and "size" in m75 if isinstance(m["counts"], list): # uncompressed RLEs76 h, w = m["size"]77 assert h == height and w == width78 m = mask_util.frPyObjects(m, h, w)79 self._mask = mask_util.decode(m)[:, :]80 return81 82 if isinstance(m, list): # list[ndarray]83 self._polygons = [np.asarray(x).reshape(-1) for x in m]84 return85 86 if isinstance(m, np.ndarray): # assumed to be a binary mask87 assert m.shape[1] != 2, m.shape88 assert m.shape == (89 height,90 width,91 ), f"mask shape: {m.shape}, target dims: {height}, {width}"92 self._mask = m.astype("uint8")93 return94 95 raise ValueError("GenericMask cannot handle object {} of type '{}'".format(m, type(m)))96 97 @property98 def mask(self):99 if self._mask is None:100 self._mask = self.polygons_to_mask(self._polygons)101 return self._mask102 103 @property104 def polygons(self):105 if self._polygons is None:106 self._polygons, self._has_holes = self.mask_to_polygons(self._mask)107 return self._polygons108 109 @property110 def has_holes(self):111 if self._has_holes is None:112 if self._mask is not None:113 self._polygons, self._has_holes = self.mask_to_polygons(self._mask)114 else:115 self._has_holes = False # if original format is polygon, does not have holes116 return self._has_holes117 118 def mask_to_polygons(self, mask):119 # cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level120 # hierarchy. External contours (boundary) of the object are placed in hierarchy-1.121 # Internal contours (holes) are placed in hierarchy-2.122 # cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours.123 mask = np.ascontiguousarray(mask) # some versions of cv2 does not support incontiguous arr124 res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)125 hierarchy = res[-1]126 if hierarchy is None: # empty mask127 return [], False128 has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0129 res = res[-2]130 res = [x.flatten() for x in res]131 # These coordinates from OpenCV are integers in range [0, W-1 or H-1].132 # We add 0.5 to turn them into real-value coordinate space. A better solution133 # would be to first +0.5 and then dilate the returned polygon by 0.5.134 res = [x + 0.5 for x in res if len(x) >= 6]135 return res, has_holes136 137 def polygons_to_mask(self, polygons):138 rle = mask_util.frPyObjects(polygons, self.height, self.width)139 rle = mask_util.merge(rle)140 return mask_util.decode(rle)[:, :]141 142 def area(self):143 return self.mask.sum()144 145 def bbox(self):146 p = mask_util.frPyObjects(self.polygons, self.height, self.width)147 p = mask_util.merge(p)148 bbox = mask_util.toBbox(p)149 bbox[2] += bbox[0]150 bbox[3] += bbox[1]151 return bbox152 153 154class _PanopticPrediction:155 """156 Unify different panoptic annotation/prediction formats157 """158 159 def __init__(self, panoptic_seg, segments_info, metadata=None):160 if segments_info is None:161 assert metadata is not None162 # If "segments_info" is None, we assume "panoptic_img" is a163 # H*W int32 image storing the panoptic_id in the format of164 # category_id * label_divisor + instance_id. We reserve -1 for165 # VOID label.166 label_divisor = metadata.label_divisor167 segments_info = []168 for panoptic_label in np.unique(panoptic_seg.numpy()):169 if panoptic_label == -1:170 # VOID region.171 continue172 pred_class = panoptic_label // label_divisor173 isthing = pred_class in metadata.thing_dataset_id_to_contiguous_id.values()174 segments_info.append(175 {176 "id": int(panoptic_label),177 "category_id": int(pred_class),178 "isthing": bool(isthing),179 }180 )181 del metadata182 183 self._seg = panoptic_seg184 185 self._sinfo = {s["id"]: s for s in segments_info} # seg id -> seg info186 segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True)187 areas = areas.numpy()188 sorted_idxs = np.argsort(-areas)189 self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs]190 self._seg_ids = self._seg_ids.tolist()191 for sid, area in zip(self._seg_ids, self._seg_areas):192 if sid in self._sinfo:193 self._sinfo[sid]["area"] = float(area)194 195 def non_empty_mask(self):196 """197 Returns:198 (H, W) array, a mask for all pixels that have a prediction199 """200 empty_ids = []201 for id in self._seg_ids:202 if id not in self._sinfo:203 empty_ids.append(id)204 if len(empty_ids) == 0:205 return np.zeros(self._seg.shape, dtype=np.uint8)206 assert (207 len(empty_ids) == 1208 ), ">1 ids corresponds to no labels. This is currently not supported"209 return (self._seg != empty_ids[0]).numpy().astype(np.bool)210 211 def semantic_masks(self):212 for sid in self._seg_ids:213 sinfo = self._sinfo.get(sid)214 if sinfo is None or sinfo["isthing"]:215 # Some pixels (e.g. id 0 in PanopticFPN) have no instance or semantic predictions.216 continue217 yield (self._seg == sid).numpy().astype(np.bool), sinfo218 219 def instance_masks(self):220 for sid in self._seg_ids:221 sinfo = self._sinfo.get(sid)222 if sinfo is None or not sinfo["isthing"]:223 continue224 mask = (self._seg == sid).numpy().astype(np.bool)225 if mask.sum() > 0:226 yield mask, sinfo227 228 229def _create_text_labels(classes, scores, class_names, is_crowd=None):230 """231 Args:232 classes (list[int] or None):233 scores (list[float] or None):234 class_names (list[str] or None):235 is_crowd (list[bool] or None):236 237 Returns:238 list[str] or None239 """240 labels = None241 if classes is not None:242 if class_names is not None and len(class_names) > 0:243 labels = [class_names[i] for i in classes]244 else:245 labels = [str(i) for i in classes]246 if scores is not None:247 if labels is None:248 labels = ["{:.0f}%".format(s * 100) for s in scores]249 else:250 labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)]251 if labels is not None and is_crowd is not None:252 labels = [l + ("|crowd" if crowd else "") for l, crowd in zip(labels, is_crowd)]253 return labels254 255 256class VisImage:257 def __init__(self, img, scale=1.0):258 """259 Args:260 img (ndarray): an RGB image of shape (H, W, 3) in range [0, 255].261 scale (float): scale the input image262 """263 self.img = img264 self.scale = scale265 self.width, self.height = img.shape[1], img.shape[0]266 self._setup_figure(img)267 268 def _setup_figure(self, img):269 """270 Args:271 Same as in :meth:`__init__()`.272 273 Returns:274 fig (matplotlib.pyplot.figure): top level container for all the image plot elements.275 ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system.276 """277 fig = mplfigure.Figure(frameon=False)278 self.dpi = fig.get_dpi()279 # add a small 1e-2 to avoid precision lost due to matplotlib's truncation280 # (https://github.com/matplotlib/matplotlib/issues/15363)281 fig.set_size_inches(282 (self.width * self.scale + 1e-2) / self.dpi,283 (self.height * self.scale + 1e-2) / self.dpi,284 )285 self.canvas = FigureCanvasAgg(fig)286 # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig)287 ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])288 ax.axis("off")289 self.fig = fig290 self.ax = ax291 self.reset_image(img)292 293 def reset_image(self, img):294 """295 Args:296 img: same as in __init__297 """298 img = img.astype("uint8")299 self.ax.imshow(img, extent=(0, self.width, self.height, 0), interpolation="nearest")300 301 def save(self, filepath):302 """303 Args:304 filepath (str): a string that contains the absolute path, including the file name, where305 the visualized image will be saved.306 """307 self.fig.savefig(filepath)308 309 def get_image(self):310 """311 Returns:312 ndarray:313 the visualized image of shape (H, W, 3) (RGB) in uint8 type.314 The shape is scaled w.r.t the input image using the given `scale` argument.315 """316 canvas = self.canvas317 s, (width, height) = canvas.print_to_buffer()318 # buf = io.BytesIO() # works for cairo backend319 # canvas.print_rgba(buf)320 # width, height = self.width, self.height321 # s = buf.getvalue()322 323 buffer = np.frombuffer(s, dtype="uint8")324 325 img_rgba = buffer.reshape(height, width, 4)326 rgb, alpha = np.split(img_rgba, [3], axis=2)327 return rgb.astype("uint8")328 329 330class Visualizer:331 """332 Visualizer that draws data about detection/segmentation on images.333 334 It contains methods like `draw_{text,box,circle,line,binary_mask,polygon}`335 that draw primitive objects to images, as well as high-level wrappers like336 `draw_{instance_predictions,sem_seg,panoptic_seg_predictions,dataset_dict}`337 that draw composite data in some pre-defined style.338 339 Note that the exact visualization style for the high-level wrappers are subject to change.340 Style such as color, opacity, label contents, visibility of labels, or even the visibility341 of objects themselves (e.g. when the object is too small) may change according342 to different heuristics, as long as the results still look visually reasonable.343 344 To obtain a consistent style, you can implement custom drawing functions with the345 abovementioned primitive methods instead. If you need more customized visualization346 styles, you can process the data yourself following their format documented in347 tutorials (:doc:`/tutorials/models`, :doc:`/tutorials/datasets`). This class does not348 intend to satisfy everyone's preference on drawing styles.349 350 This visualizer focuses on high rendering quality rather than performance. It is not351 designed to be used for real-time applications.352 """353 354 # TODO implement a fast, rasterized version using OpenCV355 356 def __init__(self, img_rgb, metadata=None, scale=1.0, instance_mode=ColorMode.IMAGE):357 """358 Args:359 img_rgb: a numpy array of shape (H, W, C), where H and W correspond to360 the height and width of the image respectively. C is the number of361 color channels. The image is required to be in RGB format since that362 is a requirement of the Matplotlib library. The image is also expected363 to be in the range [0, 255].364 metadata (Metadata): dataset metadata (e.g. class names and colors)365 instance_mode (ColorMode): defines one of the pre-defined style for drawing366 instances on an image.367 """368 self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8)369 if metadata is None:370 metadata = MetadataCatalog.get("__nonexist__")371 self.metadata = metadata372 self.output = VisImage(self.img, scale=scale)373 self.cpu_device = torch.device("cpu")374 375 # too small texts are useless, therefore clamp to 9376 self._default_font_size = max(377 np.sqrt(self.output.height * self.output.width) // 90, 10 // scale378 )379 self._default_font_size = 18380 self._instance_mode = instance_mode381 self.keypoint_threshold = _KEYPOINT_THRESHOLD382 383 def draw_instance_predictions(self, predictions):384 """385 Draw instance-level prediction results on an image.386 387 Args:388 predictions (Instances): the output of an instance detection/segmentation389 model. Following fields will be used to draw:390 "pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").391 392 Returns:393 output (VisImage): image object with visualizations.394 """395 boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None396 scores = predictions.scores if predictions.has("scores") else None397 classes = predictions.pred_classes.tolist() if predictions.has("pred_classes") else None398 labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))399 keypoints = predictions.pred_keypoints if predictions.has("pred_keypoints") else None400 401 keep = (scores > 0.8).cpu()402 boxes = boxes[keep]403 scores = scores[keep]404 classes = np.array(classes)405 classes = classes[np.array(keep)]406 labels = np.array(labels)407 labels = labels[np.array(keep)]408 409 if predictions.has("pred_masks"):410 masks = np.asarray(predictions.pred_masks)411 masks = masks[np.array(keep)]412 masks = [GenericMask(x, self.output.height, self.output.width) for x in masks]413 else:414 masks = None415 416 if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"):417 # if self.metadata.get("thing_colors"):418 colors = [419 self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in classes420 ]421 alpha = 0.4422 else:423 colors = None424 alpha = 0.4425 426 if self._instance_mode == ColorMode.IMAGE_BW:427 self.output.reset_image(428 self._create_grayscale_image(429 (predictions.pred_masks.any(dim=0) > 0).numpy()430 if predictions.has("pred_masks")431 else None432 )433 )434 alpha = 0.3435 436 self.overlay_instances(437 masks=masks,438 boxes=boxes,439 labels=labels,440 keypoints=keypoints,441 assigned_colors=colors,442 alpha=alpha,443 )444 return self.output445 446 def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.7):447 """448 Draw semantic segmentation predictions/labels.449 450 Args:451 sem_seg (Tensor or ndarray): the segmentation of shape (H, W).452 Each value is the integer label of the pixel.453 area_threshold (int): segments with less than `area_threshold` are not drawn.454 alpha (float): the larger it is, the more opaque the segmentations are.455 456 Returns:457 output (VisImage): image object with visualizations.458 """459 if isinstance(sem_seg, torch.Tensor):460 sem_seg = sem_seg.numpy()461 labels, areas = np.unique(sem_seg, return_counts=True)462 sorted_idxs = np.argsort(-areas).tolist()463 labels = labels[sorted_idxs]464 for label in filter(lambda l: l < len(self.metadata.stuff_classes), labels):465 try:466 mask_color = [x / 255 for x in self.metadata.stuff_colors[label]]467 except (AttributeError, IndexError):468 mask_color = None469 470 binary_mask = (sem_seg == label).astype(np.uint8)471 text = self.metadata.stuff_classes[label]472 self.draw_binary_mask(473 binary_mask,474 color=mask_color,475 edge_color=_OFF_WHITE,476 text=text,477 alpha=alpha,478 area_threshold=area_threshold,479 )480 return self.output481 482 def draw_panoptic_seg(self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7):483 """484 Draw panoptic prediction annotations or results.485 486 Args:487 panoptic_seg (Tensor): of shape (height, width) where the values are ids for each488 segment.489 segments_info (list[dict] or None): Describe each segment in `panoptic_seg`.490 If it is a ``list[dict]``, each dict contains keys "id", "category_id".491 If None, category id of each pixel is computed by492 ``pixel // metadata.label_divisor``.493 area_threshold (int): stuff segments with less than `area_threshold` are not drawn.494 495 Returns:496 output (VisImage): image object with visualizations.497 """498 pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)499 500 if self._instance_mode == ColorMode.IMAGE_BW:501 self.output.reset_image(self._create_grayscale_image(pred.non_empty_mask()))502 503 # draw mask for all semantic segments first i.e. "stuff"504 for mask, sinfo in pred.semantic_masks():505 category_idx = sinfo["category_id"]506 try:507 mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]508 except AttributeError:509 mask_color = None510 511 text = self.metadata.stuff_classes[category_idx]512 self.draw_binary_mask(513 mask,514 color=mask_color,515 edge_color=_OFF_WHITE,516 text=text,517 alpha=alpha,518 area_threshold=area_threshold,519 )520 521 # draw mask for all instances second522 all_instances = list(pred.instance_masks())523 if len(all_instances) == 0:524 return self.output525 masks, sinfo = list(zip(*all_instances))526 category_ids = [x["category_id"] for x in sinfo]527 528 try:529 scores = [x["score"] for x in sinfo]530 except KeyError:531 scores = None532 labels = _create_text_labels(533 category_ids, scores, self.metadata.thing_classes, [x.get("iscrowd", 0) for x in sinfo]534 )535 536 try:537 colors = [538 self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in category_ids539 ]540 except AttributeError:541 colors = None542 self.overlay_instances(masks=masks, labels=labels, assigned_colors=colors, alpha=alpha)543 544 return self.output545 546 draw_panoptic_seg_predictions = draw_panoptic_seg # backward compatibility547 548 def draw_dataset_dict(self, dic):549 """550 Draw annotations/segmentaions in Detectron2 Dataset format.551 552 Args:553 dic (dict): annotation/segmentation data of one image, in Detectron2 Dataset format.554 555 Returns:556 output (VisImage): image object with visualizations.557 """558 annos = dic.get("annotations", None)559 if annos:560 if "segmentation" in annos[0]:561 masks = [x["segmentation"] for x in annos]562 else:563 masks = None564 if "keypoints" in annos[0]:565 keypts = [x["keypoints"] for x in annos]566 keypts = np.array(keypts).reshape(len(annos), -1, 3)567 else:568 keypts = None569 570 boxes = [571 BoxMode.convert(x["bbox"], x["bbox_mode"], BoxMode.XYXY_ABS)572 if len(x["bbox"]) == 4573 else x["bbox"]574 for x in annos575 ]576 577 colors = None578 category_ids = [x["category_id"] for x in annos]579 if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"):580 colors = [581 self._jitter([x / 255 for x in self.metadata.thing_colors[c]])582 for c in category_ids583 ]584 names = self.metadata.get("thing_classes", None)585 labels = _create_text_labels(586 category_ids,587 scores=None,588 class_names=names,589 is_crowd=[x.get("iscrowd", 0) for x in annos],590 )591 self.overlay_instances(592 labels=labels, boxes=boxes, masks=masks, keypoints=keypts, assigned_colors=colors593 )594 595 sem_seg = dic.get("sem_seg", None)596 if sem_seg is None and "sem_seg_file_name" in dic:597 with PathManager.open(dic["sem_seg_file_name"], "rb") as f:598 sem_seg = Image.open(f)599 sem_seg = np.asarray(sem_seg, dtype="uint8")600 if sem_seg is not None:601 self.draw_sem_seg(sem_seg, area_threshold=0, alpha=0.4)602 603 pan_seg = dic.get("pan_seg", None)604 if pan_seg is None and "pan_seg_file_name" in dic:605 with PathManager.open(dic["pan_seg_file_name"], "rb") as f:606 pan_seg = Image.open(f)607 pan_seg = np.asarray(pan_seg)608 from panopticapi.utils import rgb2id609 610 pan_seg = rgb2id(pan_seg)611 if pan_seg is not None:612 segments_info = dic["segments_info"]613 pan_seg = torch.tensor(pan_seg)614 self.draw_panoptic_seg(pan_seg, segments_info, area_threshold=0, alpha=0.7)615 return self.output616 617 def overlay_instances(618 self,619 *,620 boxes=None,621 labels=None,622 masks=None,623 keypoints=None,624 assigned_colors=None,625 alpha=0.5,626 ):627 """628 Args:629 boxes (Boxes, RotatedBoxes or ndarray): either a :class:`Boxes`,630 or an Nx4 numpy array of XYXY_ABS format for the N objects in a single image,631 or a :class:`RotatedBoxes`,632 or an Nx5 numpy array of (x_center, y_center, width, height, angle_degrees) format633 for the N objects in a single image,634 labels (list[str]): the text to be displayed for each instance.635 masks (masks-like object): Supported types are:636 637 * :class:`detectron2.structures.PolygonMasks`,638 :class:`detectron2.structures.BitMasks`.639 * list[list[ndarray]]: contains the segmentation masks for all objects in one image.640 The first level of the list corresponds to individual instances. The second641 level to all the polygon that compose the instance, and the third level642 to the polygon coordinates. The third level should have the format of643 [x0, y0, x1, y1, ..., xn, yn] (n >= 3).644 * list[ndarray]: each ndarray is a binary mask of shape (H, W).645 * list[dict]: each dict is a COCO-style RLE.646 keypoints (Keypoint or array like): an array-like object of shape (N, K, 3),647 where the N is the number of instances and K is the number of keypoints.648 The last dimension corresponds to (x, y, visibility or score).649 assigned_colors (list[matplotlib.colors]): a list of colors, where each color650 corresponds to each mask or box in the image. Refer to 'matplotlib.colors'651 for full list of formats that the colors are accepted in.652 Returns:653 output (VisImage): image object with visualizations.654 """655 num_instances = 0656 if boxes is not None:657 boxes = self._convert_boxes(boxes)658 num_instances = len(boxes)659 if masks is not None:660 masks = self._convert_masks(masks)661 if num_instances:662 assert len(masks) == num_instances663 else:664 num_instances = len(masks)665 if keypoints is not None:666 if num_instances:667 assert len(keypoints) == num_instances668 else:669 num_instances = len(keypoints)670 keypoints = self._convert_keypoints(keypoints)671 if labels is not None:672 assert len(labels) == num_instances673 if assigned_colors is None:674 assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)]675 if num_instances == 0:676 return self.output677 if boxes is not None and boxes.shape[1] == 5:678 return self.overlay_rotated_instances(679 boxes=boxes, labels=labels, assigned_colors=assigned_colors680 )681 682 # Display in largest to smallest order to reduce occlusion.683 areas = None684 if boxes is not None:685 areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1)686 elif masks is not None:687 areas = np.asarray([x.area() for x in masks])688 689 if areas is not None:690 sorted_idxs = np.argsort(-areas).tolist()691 # Re-order overlapped instances in descending order.692 boxes = boxes[sorted_idxs] if boxes is not None else None693 labels = [labels[k] for k in sorted_idxs] if labels is not None else None694 masks = [masks[idx] for idx in sorted_idxs] if masks is not None else None695 assigned_colors = [assigned_colors[idx] for idx in sorted_idxs]696 keypoints = keypoints[sorted_idxs] if keypoints is not None else None697 698 for i in range(num_instances):699 color = assigned_colors[i]700 if boxes is not None:701 self.draw_box(boxes[i], edge_color=color)702 703 if masks is not None:704 for segment in masks[i].polygons:705 self.draw_polygon(segment.reshape(-1, 2), color, alpha=alpha)706 707 if labels is not None:708 # first get a box709 if boxes is not None:710 x0, y0, x1, y1 = boxes[i]711 text_pos = (x0, y0) # if drawing boxes, put text on the box corner.712 horiz_align = "left"713 elif masks is not None:714 # skip small mask without polygon715 if len(masks[i].polygons) == 0:716 continue717 718 x0, y0, x1, y1 = masks[i].bbox()719 720 # draw text in the center (defined by median) when box is not drawn721 # median is less sensitive to outliers.722 text_pos = np.median(masks[i].mask.nonzero(), axis=1)[::-1]723 horiz_align = "center"724 else:725 continue # drawing the box confidence for keypoints isn't very useful.726 # for small objects, draw text at the side to avoid occlusion727 instance_area = (y1 - y0) * (x1 - x0)728 if (729 instance_area < _SMALL_OBJECT_AREA_THRESH * self.output.scale730 or y1 - y0 < 40 * self.output.scale731 ):732 if y1 >= self.output.height - 5:733 text_pos = (x1, y0)734 else:735 text_pos = (x0, y1)736 737 height_ratio = (y1 - y0) / np.sqrt(self.output.height * self.output.width)738 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)739 font_size = (740 np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2)741 * 0.5742 * self._default_font_size743 )744 self.draw_text(745 labels[i],746 text_pos,747 color=lighter_color,748 horizontal_alignment=horiz_align,749 font_size=font_size,750 )751 752 # draw keypoints753 if keypoints is not None:754 for keypoints_per_instance in keypoints:755 self.draw_and_connect_keypoints(keypoints_per_instance)756 757 return self.output758 759 def overlay_rotated_instances(self, boxes=None, labels=None, assigned_colors=None):760 """761 Args:762 boxes (ndarray): an Nx5 numpy array of763 (x_center, y_center, width, height, angle_degrees) format764 for the N objects in a single image.765 labels (list[str]): the text to be displayed for each instance.766 assigned_colors (list[matplotlib.colors]): a list of colors, where each color767 corresponds to each mask or box in the image. Refer to 'matplotlib.colors'768 for full list of formats that the colors are accepted in.769 770 Returns:771 output (VisImage): image object with visualizations.772 """773 num_instances = len(boxes)774 775 if assigned_colors is None:776 assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)]777 if num_instances == 0:778 return self.output779 780 # Display in largest to smallest order to reduce occlusion.781 if boxes is not None:782 areas = boxes[:, 2] * boxes[:, 3]783 784 sorted_idxs = np.argsort(-areas).tolist()785 # Re-order overlapped instances in descending order.786 boxes = boxes[sorted_idxs]787 labels = [labels[k] for k in sorted_idxs] if labels is not None else None788 colors = [assigned_colors[idx] for idx in sorted_idxs]789 790 for i in range(num_instances):791 self.draw_rotated_box_with_label(792 boxes[i], edge_color=colors[i], label=labels[i] if labels is not None else None793 )794 795 return self.output796 797 def draw_and_connect_keypoints(self, keypoints):798 """799 Draws keypoints of an instance and follows the rules for keypoint connections800 to draw lines between appropriate keypoints. This follows color heuristics for801 line color.802 803 Args:804 keypoints (Tensor): a tensor of shape (K, 3), where K is the number of keypoints805 and the last dimension corresponds to (x, y, probability).806 807 Returns:808 output (VisImage): image object with visualizations.809 """810 visible = {}811 keypoint_names = self.metadata.get("keypoint_names")812 for idx, keypoint in enumerate(keypoints):813 814 # draw keypoint815 x, y, prob = keypoint816 if prob > self.keypoint_threshold:817 self.draw_circle((x, y), color=_RED)818 if keypoint_names:819 keypoint_name = keypoint_names[idx]820 visible[keypoint_name] = (x, y)821 822 if self.metadata.get("keypoint_connection_rules"):823 for kp0, kp1, color in self.metadata.keypoint_connection_rules:824 if kp0 in visible and kp1 in visible:825 x0, y0 = visible[kp0]826 x1, y1 = visible[kp1]827 color = tuple(x / 255.0 for x in color)828 self.draw_line([x0, x1], [y0, y1], color=color)829 830 # draw lines from nose to mid-shoulder and mid-shoulder to mid-hip831 # Note that this strategy is specific to person keypoints.832 # For other keypoints, it should just do nothing833 try:834 ls_x, ls_y = visible["left_shoulder"]835 rs_x, rs_y = visible["right_shoulder"]836 mid_shoulder_x, mid_shoulder_y = (ls_x + rs_x) / 2, (ls_y + rs_y) / 2837 except KeyError:838 pass839 else:840 # draw line from nose to mid-shoulder841 nose_x, nose_y = visible.get("nose", (None, None))842 if nose_x is not None:843 self.draw_line([nose_x, mid_shoulder_x], [nose_y, mid_shoulder_y], color=_RED)844 845 try:846 # draw line from mid-shoulder to mid-hip847 lh_x, lh_y = visible["left_hip"]848 rh_x, rh_y = visible["right_hip"]849 except KeyError:850 pass851 else:852 mid_hip_x, mid_hip_y = (lh_x + rh_x) / 2, (lh_y + rh_y) / 2853 self.draw_line([mid_hip_x, mid_shoulder_x], [mid_hip_y, mid_shoulder_y], color=_RED)854 return self.output855 856 """857 Primitive drawing functions:858 """859 860 def draw_text(861 self,862 text,863 position,864 *,865 font_size=None,866 color="g",867 horizontal_alignment="center",868 rotation=0,869 ):870 """871 Args:872 text (str): class label873 position (tuple): a tuple of the x and y coordinates to place text on image.874 font_size (int, optional): font of the text. If not provided, a font size875 proportional to the image width is calculated and used.876 color: color of the text. Refer to `matplotlib.colors` for full list877 of formats that are accepted.878 horizontal_alignment (str): see `matplotlib.text.Text`879 rotation: rotation angle in degrees CCW880 881 Returns:882 output (VisImage): image object with text drawn.883 """884 if not font_size:885 font_size = self._default_font_size886 887 # since the text background is dark, we don't want the text to be dark888 color = np.maximum(list(mplc.to_rgb(color)), 0.2)889 color[np.argmax(color)] = max(0.8, np.max(color))890 891 x, y = position892 self.output.ax.text(893 x,894 y,895 text,896 size=font_size * self.output.scale,897 family="sans-serif",898 bbox={"facecolor": "black", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"},899 verticalalignment="top",900 horizontalalignment=horizontal_alignment,901 color=color,902 zorder=10,903 rotation=rotation,904 )905 return self.output906 907 def draw_box(self, box_coord, alpha=0.5, edge_color="g", line_style="-"):908 """909 Args:910 box_coord (tuple): a tuple containing x0, y0, x1, y1 coordinates, where x0 and y0911 are the coordinates of the image's top left corner. x1 and y1 are the912 coordinates of the image's bottom right corner.913 alpha (float): blending efficient. Smaller values lead to more transparent masks.914 edge_color: color of the outline of the box. Refer to `matplotlib.colors`915 for full list of formats that are accepted.916 line_style (string): the string to use to create the outline of the boxes.917 918 Returns:919 output (VisImage): image object with box drawn.920 """921 x0, y0, x1, y1 = box_coord922 width = x1 - x0923 height = y1 - y0924 925 linewidth = max(self._default_font_size / 4, 1)926 927 self.output.ax.add_patch(928 mpl.patches.Rectangle(929 (x0, y0),930 width,931 height,932 fill=False,933 edgecolor=edge_color,934 linewidth=linewidth * self.output.scale,935 alpha=alpha,936 linestyle=line_style,937 )938 )939 return self.output940 941 def draw_rotated_box_with_label(942 self, rotated_box, alpha=0.5, edge_color="g", line_style="-", label=None943 ):944 """945 Draw a rotated box with label on its top-left corner.946 947 Args:948 rotated_box (tuple): a tuple containing (cnt_x, cnt_y, w, h, angle),949 where cnt_x and cnt_y are the center coordinates of the box.950 w and h are the width and height of the box. angle represents how951 many degrees the box is rotated CCW with regard to the 0-degree box.952 alpha (float): blending efficient. Smaller values lead to more transparent masks.953 edge_color: color of the outline of the box. Refer to `matplotlib.colors`954 for full list of formats that are accepted.955 line_style (string): the string to use to create the outline of the boxes.956 label (string): label for rotated box. It will not be rendered when set to None.957 958 Returns:959 output (VisImage): image object with box drawn.960 """961 cnt_x, cnt_y, w, h, angle = rotated_box962 area = w * h963 # use thinner lines when the box is small964 linewidth = self._default_font_size / (965 6 if area < _SMALL_OBJECT_AREA_THRESH * self.output.scale else 3966 )967 968 theta = angle * math.pi / 180.0969 c = math.cos(theta)970 s = math.sin(theta)971 rect = [(-w / 2, h / 2), (-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2)]972 # x: left->right ; y: top->down973 rotated_rect = [(s * yy + c * xx + cnt_x, c * yy - s * xx + cnt_y) for (xx, yy) in rect]974 for k in range(4):975 j = (k + 1) % 4976 self.draw_line(977 [rotated_rect[k][0], rotated_rect[j][0]],978 [rotated_rect[k][1], rotated_rect[j][1]],979 color=edge_color,980 linestyle="--" if k == 1 else line_style,981 linewidth=linewidth,982 )983 984 if label is not None:985 text_pos = rotated_rect[1] # topleft corner986 987 height_ratio = h / np.sqrt(self.output.height * self.output.width)988 label_color = self._change_color_brightness(edge_color, brightness_factor=0.7)989 font_size = (990 np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) * 0.5 * self._default_font_size991 )992 self.draw_text(label, text_pos, color=label_color, font_size=font_size, rotation=angle)993 994 return self.output995 996 def draw_circle(self, circle_coord, color, radius=3):997 """998 Args:999 circle_coord (list(int) or tuple(int)): contains the x and y coordinates1000 of the center of the circle.1001 color: color of the polygon. Refer to `matplotlib.colors` for a full list of1002 formats that are accepted.1003 radius (int): radius of the circle.1004 1005 Returns:1006 output (VisImage): image object with box drawn.1007 """1008 x, y = circle_coord1009 self.output.ax.add_patch(1010 mpl.patches.Circle(circle_coord, radius=radius, fill=True, color=color)1011 )1012 return self.output1013 1014 def draw_line(self, x_data, y_data, color, linestyle="-", linewidth=None):1015 """1016 Args:1017 x_data (list[int]): a list containing x values of all the points being drawn.1018 Length of list should match the length of y_data.1019 y_data (list[int]): a list containing y values of all the points being drawn.1020 Length of list should match the length of x_data.1021 color: color of the line. Refer to `matplotlib.colors` for a full list of1022 formats that are accepted.1023 linestyle: style of the line. Refer to `matplotlib.lines.Line2D`1024 for a full list of formats that are accepted.1025 linewidth (float or None): width of the line. When it's None,1026 a default value will be computed and used.1027 1028 Returns:1029 output (VisImage): image object with line drawn.1030 """1031 if linewidth is None:1032 linewidth = self._default_font_size / 31033 linewidth = max(linewidth, 1)1034 self.output.ax.add_line(1035 mpl.lines.Line2D(1036 x_data,1037 y_data,1038 linewidth=linewidth * self.output.scale,1039 color=color,1040 linestyle=linestyle,1041 )1042 )1043 return self.output1044 1045 def draw_binary_mask(1046 self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.7, area_threshold=101047 ):1048 """1049 Args:1050 binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and1051 W is the image width. Each value in the array is either a 0 or 1 value of uint81052 type.1053 color: color of the mask. Refer to `matplotlib.colors` for a full list of1054 formats that are accepted. If None, will pick a random color.1055 edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a1056 full list of formats that are accepted.1057 text (str): if None, will be drawn on the object1058 alpha (float): blending efficient. Smaller values lead to more transparent masks.1059 area_threshold (float): a connected component smaller than this area will not be shown.1060 1061 Returns:1062 output (VisImage): image object with mask drawn.1063 """1064 if color is None:1065 color = random_color(rgb=True, maximum=1)1066 color = mplc.to_rgb(color)1067 1068 has_valid_segment = False1069 binary_mask = binary_mask.astype("uint8") # opencv needs uint81070 mask = GenericMask(binary_mask, self.output.height, self.output.width)1071 shape2d = (binary_mask.shape[0], binary_mask.shape[1])1072 1073 if not mask.has_holes:1074 # draw polygons for regular masks1075 for segment in mask.polygons:1076 area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1]))1077 if area < (area_threshold or 0):1078 continue1079 has_valid_segment = True1080 segment = segment.reshape(-1, 2)1081 self.draw_polygon(segment, color=color, edge_color=edge_color, alpha=alpha)1082 else:1083 # TODO: Use Path/PathPatch to draw vector graphics:1084 # https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon1085 rgba = np.zeros(shape2d + (4,), dtype="float32")1086 rgba[:, :, :3] = color1087 rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha1088 has_valid_segment = True1089 self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0))1090 1091 if text is not None and has_valid_segment:1092 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)1093 self._draw_text_in_mask(binary_mask, text, lighter_color)1094 return self.output1095 1096 def draw_soft_mask(self, soft_mask, color=None, *, text=None, alpha=0.5):1097 """1098 Args:1099 soft_mask (ndarray): float array of shape (H, W), each value in [0, 1].1100 color: color of the mask. Refer to `matplotlib.colors` for a full list of1101 formats that are accepted. If None, will pick a random color.1102 text (str): if None, will be drawn on the object1103 alpha (float): blending efficient. Smaller values lead to more transparent masks.1104 1105 Returns:1106 output (VisImage): image object with mask drawn.1107 """1108 if color is None:1109 color = random_color(rgb=True, maximum=1)1110 color = mplc.to_rgb(color)1111 1112 shape2d = (soft_mask.shape[0], soft_mask.shape[1])1113 rgba = np.zeros(shape2d + (4,), dtype="float32")1114 rgba[:, :, :3] = color1115 rgba[:, :, 3] = soft_mask * alpha1116 self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0))1117 1118 if text is not None:1119 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)1120 binary_mask = (soft_mask > 0.5).astype("uint8")1121 self._draw_text_in_mask(binary_mask, text, lighter_color)1122 return self.output1123 1124 def draw_polygon(self, segment, color, edge_color=None, alpha=0.5):1125 """1126 Args:1127 segment: numpy array of shape Nx2, containing all the points in the polygon.1128 color: color of the polygon. Refer to `matplotlib.colors` for a full list of1129 formats that are accepted.1130 edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a1131 full list of formats that are accepted. If not provided, a darker shade1132 of the polygon color will be used instead.1133 alpha (float): blending efficient. Smaller values lead to more transparent masks.1134 1135 Returns:1136 output (VisImage): image object with polygon drawn.1137 """1138 if edge_color is None:1139 # make edge color darker than the polygon color1140 if alpha > 0.8:1141 edge_color = self._change_color_brightness(color, brightness_factor=-0.7)1142 else:1143 edge_color = color1144 edge_color = mplc.to_rgb(edge_color) + (1,)1145 1146 polygon = mpl.patches.Polygon(1147 segment,1148 fill=True,1149 facecolor=mplc.to_rgb(color) + (alpha,),1150 edgecolor=edge_color,1151 linewidth=max(self._default_font_size // 15 * self.output.scale, 1),1152 )1153 self.output.ax.add_patch(polygon)1154 return self.output1155 1156 """1157 Internal methods:1158 """1159 1160 def _jitter(self, color):1161 """1162 Randomly modifies given color to produce a slightly different color than the color given.1163 1164 Args:1165 color (tuple[double]): a tuple of 3 elements, containing the RGB values of the color1166 picked. The values in the list are in the [0.0, 1.0] range.1167 1168 Returns:1169 jittered_color (tuple[double]): a tuple of 3 elements, containing the RGB values of the1170 color after being jittered. The values in the list are in the [0.0, 1.0] range.1171 """1172 color = mplc.to_rgb(color)1173 # np.random.seed(0)1174 vec = np.random.rand(3)1175 # better to do it in another color space1176 vec = vec / np.linalg.norm(vec) * 0.51177 res = np.clip(vec + color, 0, 1)1178 return tuple(res)1179 1180 def _create_grayscale_image(self, mask=None):1181 """1182 Create a grayscale version of the original image.1183 The colors in masked area, if given, will be kept.1184 """1185 img_bw = self.img.astype("f4").mean(axis=2)1186 img_bw = np.stack([img_bw] * 3, axis=2)1187 if mask is not None:1188 img_bw[mask] = self.img[mask]1189 return img_bw1190 1191 def _change_color_brightness(self, color, brightness_factor):1192 """1193 Depending on the brightness_factor, gives a lighter or darker color i.e. a color with1194 less or more saturation than the original color.1195 1196 Args:1197 color: color of the polygon. Refer to `matplotlib.colors` for a full list of1198 formats that are accepted.1199 brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of1200 0 will correspond to no change, a factor in [-1.0, 0) range will result in