jawahar-konathala/Tryon2
0
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 .colormap import random_color21 22logger = logging.getLogger(__name__)23 24__all__ = ["ColorMode", "VisImage", "Visualizer"]25 26 27_SMALL_OBJECT_AREA_THRESH = 100028_LARGE_MASK_AREA_THRESH = 12000029_OFF_WHITE = (1.0, 1.0, 240.0 / 255)30_BLACK = (0, 0, 0)31_RED = (1.0, 0, 0)32 33_KEYPOINT_THRESHOLD = 0.0534 35 36@unique37class ColorMode(Enum):38 """39 Enum of different color modes to use for instance visualizations.40 """41 42 IMAGE = 043 """44 Picks a random color for every instance and overlay segmentations with low opacity.45 """46 SEGMENTATION = 147 """48 Let instances of the same category have similar colors49 (from metadata.thing_colors), and overlay them with50 high opacity. This provides more attention on the quality of segmentation.51 """52 IMAGE_BW = 253 """54 Same as IMAGE, but convert all areas without masks to gray-scale.55 Only available for drawing per-instance mask predictions.56 """57 58 59class GenericMask:60 """61 Attribute:62 polygons (list[ndarray]): list[ndarray]: polygons for this mask.63 Each ndarray has format [x, y, x, y, ...]64 mask (ndarray): a binary mask65 """66 67 def __init__(self, mask_or_polygons, height, width):68 self._mask = self._polygons = self._has_holes = None69 self.height = height70 self.width = width71 72 m = mask_or_polygons73 if isinstance(m, dict):74 # RLEs75 assert "counts" in m and "size" in m76 if isinstance(m["counts"], list): # uncompressed RLEs77 h, w = m["size"]78 assert h == height and w == width79 m = mask_util.frPyObjects(m, h, w)80 self._mask = mask_util.decode(m)[:, :]81 return82 83 if isinstance(m, list): # list[ndarray]84 self._polygons = [np.asarray(x).reshape(-1) for x in m]85 return86 87 if isinstance(m, np.ndarray): # assumed to be a binary mask88 assert m.shape[1] != 2, m.shape89 assert m.shape == (90 height,91 width,92 ), f"mask shape: {m.shape}, target dims: {height}, {width}"93 self._mask = m.astype("uint8")94 return95 96 raise ValueError("GenericMask cannot handle object {} of type '{}'".format(m, type(m)))97 98 @property99 def mask(self):100 if self._mask is None:101 self._mask = self.polygons_to_mask(self._polygons)102 return self._mask103 104 @property105 def polygons(self):106 if self._polygons is None:107 self._polygons, self._has_holes = self.mask_to_polygons(self._mask)108 return self._polygons109 110 @property111 def has_holes(self):112 if self._has_holes is None:113 if self._mask is not None:114 self._polygons, self._has_holes = self.mask_to_polygons(self._mask)115 else:116 self._has_holes = False # if original format is polygon, does not have holes117 return self._has_holes118 119 def mask_to_polygons(self, mask):120 # cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level121 # hierarchy. External contours (boundary) of the object are placed in hierarchy-1.122 # Internal contours (holes) are placed in hierarchy-2.123 # cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours.124 mask = np.ascontiguousarray(mask) # some versions of cv2 does not support incontiguous arr125 res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)126 hierarchy = res[-1]127 if hierarchy is None: # empty mask128 return [], False129 has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0130 res = res[-2]131 res = [x.flatten() for x in res]132 # These coordinates from OpenCV are integers in range [0, W-1 or H-1].133 # We add 0.5 to turn them into real-value coordinate space. A better solution134 # would be to first +0.5 and then dilate the returned polygon by 0.5.135 res = [x + 0.5 for x in res if len(x) >= 6]136 return res, has_holes137 138 def polygons_to_mask(self, polygons):139 rle = mask_util.frPyObjects(polygons, self.height, self.width)140 rle = mask_util.merge(rle)141 return mask_util.decode(rle)[:, :]142 143 def area(self):144 return self.mask.sum()145 146 def bbox(self):147 p = mask_util.frPyObjects(self.polygons, self.height, self.width)148 p = mask_util.merge(p)149 bbox = mask_util.toBbox(p)150 bbox[2] += bbox[0]151 bbox[3] += bbox[1]152 return bbox153 154 155class _PanopticPrediction:156 """157 Unify different panoptic annotation/prediction formats158 """159 160 def __init__(self, panoptic_seg, segments_info, metadata=None):161 if segments_info is None:162 assert metadata is not None163 # If "segments_info" is None, we assume "panoptic_img" is a164 # H*W int32 image storing the panoptic_id in the format of165 # category_id * label_divisor + instance_id. We reserve -1 for166 # VOID label.167 label_divisor = metadata.label_divisor168 segments_info = []169 for panoptic_label in np.unique(panoptic_seg.numpy()):170 if panoptic_label == -1:171 # VOID region.172 continue173 pred_class = panoptic_label // label_divisor174 isthing = pred_class in metadata.thing_dataset_id_to_contiguous_id.values()175 segments_info.append(176 {177 "id": int(panoptic_label),178 "category_id": int(pred_class),179 "isthing": bool(isthing),180 }181 )182 del metadata183 184 self._seg = panoptic_seg185 186 self._sinfo = {s["id"]: s for s in segments_info} # seg id -> seg info187 segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True)188 areas = areas.numpy()189 sorted_idxs = np.argsort(-areas)190 self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs]191 self._seg_ids = self._seg_ids.tolist()192 for sid, area in zip(self._seg_ids, self._seg_areas):193 if sid in self._sinfo:194 self._sinfo[sid]["area"] = float(area)195 196 def non_empty_mask(self):197 """198 Returns:199 (H, W) array, a mask for all pixels that have a prediction200 """201 empty_ids = []202 for id in self._seg_ids:203 if id not in self._sinfo:204 empty_ids.append(id)205 if len(empty_ids) == 0:206 return np.zeros(self._seg.shape, dtype=np.uint8)207 assert (208 len(empty_ids) == 1209 ), ">1 ids corresponds to no labels. This is currently not supported"210 return (self._seg != empty_ids[0]).numpy().astype(bool)211 212 def semantic_masks(self):213 for sid in self._seg_ids:214 sinfo = self._sinfo.get(sid)215 if sinfo is None or sinfo["isthing"]:216 # Some pixels (e.g. id 0 in PanopticFPN) have no instance or semantic predictions.217 continue218 yield (self._seg == sid).numpy().astype(bool), sinfo219 220 def instance_masks(self):221 for sid in self._seg_ids:222 sinfo = self._sinfo.get(sid)223 if sinfo is None or not sinfo["isthing"]:224 continue225 mask = (self._seg == sid).numpy().astype(bool)226 if mask.sum() > 0:227 yield mask, sinfo228 229 230def _create_text_labels(classes, scores, class_names, is_crowd=None):231 """232 Args:233 classes (list[int] or None):234 scores (list[float] or None):235 class_names (list[str] or None):236 is_crowd (list[bool] or None):237 238 Returns:239 list[str] or None240 """241 labels = None242 if classes is not None:243 if class_names is not None and len(class_names) > 0:244 labels = [class_names[i] for i in classes]245 else:246 labels = [str(i) for i in classes]247 if scores is not None:248 if labels is None:249 labels = ["{:.0f}%".format(s * 100) for s in scores]250 else:251 labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)]252 if labels is not None and is_crowd is not None:253 labels = [l + ("|crowd" if crowd else "") for l, crowd in zip(labels, is_crowd)]254 return labels255 256 257class VisImage:258 def __init__(self, img, scale=1.0):259 """260 Args:261 img (ndarray): an RGB image of shape (H, W, 3) in range [0, 255].262 scale (float): scale the input image263 """264 self.img = img265 self.scale = scale266 self.width, self.height = img.shape[1], img.shape[0]267 self._setup_figure(img)268 269 def _setup_figure(self, img):270 """271 Args:272 Same as in :meth:`__init__()`.273 274 Returns:275 fig (matplotlib.pyplot.figure): top level container for all the image plot elements.276 ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system.277 """278 fig = mplfigure.Figure(frameon=False)279 self.dpi = fig.get_dpi()280 # add a small 1e-2 to avoid precision lost due to matplotlib's truncation281 # (https://github.com/matplotlib/matplotlib/issues/15363)282 fig.set_size_inches(283 (self.width * self.scale + 1e-2) / self.dpi,284 (self.height * self.scale + 1e-2) / self.dpi,285 )286 self.canvas = FigureCanvasAgg(fig)287 # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig)288 ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])289 ax.axis("off")290 self.fig = fig291 self.ax = ax292 self.reset_image(img)293 294 def reset_image(self, img):295 """296 Args:297 img: same as in __init__298 """299 img = img.astype("uint8")300 self.ax.imshow(img, extent=(0, self.width, self.height, 0), interpolation="nearest")301 302 def save(self, filepath):303 """304 Args:305 filepath (str): a string that contains the absolute path, including the file name, where306 the visualized image will be saved.307 """308 self.fig.savefig(filepath)309 310 def get_image(self):311 """312 Returns:313 ndarray:314 the visualized image of shape (H, W, 3) (RGB) in uint8 type.315 The shape is scaled w.r.t the input image using the given `scale` argument.316 """317 canvas = self.canvas318 s, (width, height) = canvas.print_to_buffer()319 # buf = io.BytesIO() # works for cairo backend320 # canvas.print_rgba(buf)321 # width, height = self.width, self.height322 # s = buf.getvalue()323 324 buffer = np.frombuffer(s, dtype="uint8")325 326 img_rgba = buffer.reshape(height, width, 4)327 rgb, alpha = np.split(img_rgba, [3], axis=2)328 return rgb.astype("uint8")329 330 331class Visualizer:332 """333 Visualizer that draws data about detection/segmentation on images.334 335 It contains methods like `draw_{text,box,circle,line,binary_mask,polygon}`336 that draw primitive objects to images, as well as high-level wrappers like337 `draw_{instance_predictions,sem_seg,panoptic_seg_predictions,dataset_dict}`338 that draw composite data in some pre-defined style.339 340 Note that the exact visualization style for the high-level wrappers are subject to change.341 Style such as color, opacity, label contents, visibility of labels, or even the visibility342 of objects themselves (e.g. when the object is too small) may change according343 to different heuristics, as long as the results still look visually reasonable.344 345 To obtain a consistent style, you can implement custom drawing functions with the346 abovementioned primitive methods instead. If you need more customized visualization347 styles, you can process the data yourself following their format documented in348 tutorials (:doc:`/tutorials/models`, :doc:`/tutorials/datasets`). This class does not349 intend to satisfy everyone's preference on drawing styles.350 351 This visualizer focuses on high rendering quality rather than performance. It is not352 designed to be used for real-time applications.353 """354 355 # TODO implement a fast, rasterized version using OpenCV356 357 def __init__(self, img_rgb, metadata=None, scale=1.0, instance_mode=ColorMode.IMAGE):358 """359 Args:360 img_rgb: a numpy array of shape (H, W, C), where H and W correspond to361 the height and width of the image respectively. C is the number of362 color channels. The image is required to be in RGB format since that363 is a requirement of the Matplotlib library. The image is also expected364 to be in the range [0, 255].365 metadata (Metadata): dataset metadata (e.g. class names and colors)366 instance_mode (ColorMode): defines one of the pre-defined style for drawing367 instances on an image.368 """369 self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8)370 if metadata is None:371 metadata = MetadataCatalog.get("__nonexist__")372 self.metadata = metadata373 self.output = VisImage(self.img, scale=scale)374 self.cpu_device = torch.device("cpu")375 376 # too small texts are useless, therefore clamp to 9377 self._default_font_size = max(378 np.sqrt(self.output.height * self.output.width) // 90, 10 // scale379 )380 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 if predictions.has("pred_masks"):402 masks = np.asarray(predictions.pred_masks)403 masks = [GenericMask(x, self.output.height, self.output.width) for x in masks]404 else:405 masks = None406 407 if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"):408 colors = [409 self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in classes410 ]411 alpha = 0.8412 else:413 colors = None414 alpha = 0.5415 416 if self._instance_mode == ColorMode.IMAGE_BW:417 self.output.reset_image(418 self._create_grayscale_image(419 (predictions.pred_masks.any(dim=0) > 0).numpy()420 if predictions.has("pred_masks")421 else None422 )423 )424 alpha = 0.3425 426 self.overlay_instances(427 masks=masks,428 boxes=boxes,429 labels=labels,430 keypoints=keypoints,431 assigned_colors=colors,432 alpha=alpha,433 )434 return self.output435 436 def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8):437 """438 Draw semantic segmentation predictions/labels.439 440 Args:441 sem_seg (Tensor or ndarray): the segmentation of shape (H, W).442 Each value is the integer label of the pixel.443 area_threshold (int): segments with less than `area_threshold` are not drawn.444 alpha (float): the larger it is, the more opaque the segmentations are.445 446 Returns:447 output (VisImage): image object with visualizations.448 """449 if isinstance(sem_seg, torch.Tensor):450 sem_seg = sem_seg.numpy()451 labels, areas = np.unique(sem_seg, return_counts=True)452 sorted_idxs = np.argsort(-areas).tolist()453 labels = labels[sorted_idxs]454 for label in filter(lambda l: l < len(self.metadata.stuff_classes), labels):455 try:456 mask_color = [x / 255 for x in self.metadata.stuff_colors[label]]457 except (AttributeError, IndexError):458 mask_color = None459 460 binary_mask = (sem_seg == label).astype(np.uint8)461 text = self.metadata.stuff_classes[label]462 self.draw_binary_mask(463 binary_mask,464 color=mask_color,465 edge_color=_OFF_WHITE,466 text=text,467 alpha=alpha,468 area_threshold=area_threshold,469 )470 return self.output471 472 def draw_panoptic_seg(self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7):473 """474 Draw panoptic prediction annotations or results.475 476 Args:477 panoptic_seg (Tensor): of shape (height, width) where the values are ids for each478 segment.479 segments_info (list[dict] or None): Describe each segment in `panoptic_seg`.480 If it is a ``list[dict]``, each dict contains keys "id", "category_id".481 If None, category id of each pixel is computed by482 ``pixel // metadata.label_divisor``.483 area_threshold (int): stuff segments with less than `area_threshold` are not drawn.484 485 Returns:486 output (VisImage): image object with visualizations.487 """488 pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)489 490 if self._instance_mode == ColorMode.IMAGE_BW:491 self.output.reset_image(self._create_grayscale_image(pred.non_empty_mask()))492 493 # draw mask for all semantic segments first i.e. "stuff"494 for mask, sinfo in pred.semantic_masks():495 category_idx = sinfo["category_id"]496 try:497 mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]498 except AttributeError:499 mask_color = None500 501 text = self.metadata.stuff_classes[category_idx]502 self.draw_binary_mask(503 mask,504 color=mask_color,505 edge_color=_OFF_WHITE,506 text=text,507 alpha=alpha,508 area_threshold=area_threshold,509 )510 511 # draw mask for all instances second512 all_instances = list(pred.instance_masks())513 if len(all_instances) == 0:514 return self.output515 masks, sinfo = list(zip(*all_instances))516 category_ids = [x["category_id"] for x in sinfo]517 518 try:519 scores = [x["score"] for x in sinfo]520 except KeyError:521 scores = None522 labels = _create_text_labels(523 category_ids, scores, self.metadata.thing_classes, [x.get("iscrowd", 0) for x in sinfo]524 )525 526 try:527 colors = [528 self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in category_ids529 ]530 except AttributeError:531 colors = None532 self.overlay_instances(masks=masks, labels=labels, assigned_colors=colors, alpha=alpha)533 534 return self.output535 536 draw_panoptic_seg_predictions = draw_panoptic_seg # backward compatibility537 538 def draw_dataset_dict(self, dic):539 """540 Draw annotations/segmentations in Detectron2 Dataset format.541 542 Args:543 dic (dict): annotation/segmentation data of one image, in Detectron2 Dataset format.544 545 Returns:546 output (VisImage): image object with visualizations.547 """548 annos = dic.get("annotations", None)549 if annos:550 if "segmentation" in annos[0]:551 masks = [x["segmentation"] for x in annos]552 else:553 masks = None554 if "keypoints" in annos[0]:555 keypts = [x["keypoints"] for x in annos]556 keypts = np.array(keypts).reshape(len(annos), -1, 3)557 else:558 keypts = None559 560 boxes = [561 BoxMode.convert(x["bbox"], x["bbox_mode"], BoxMode.XYXY_ABS)562 if len(x["bbox"]) == 4563 else x["bbox"]564 for x in annos565 ]566 567 colors = None568 category_ids = [x["category_id"] for x in annos]569 if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"):570 colors = [571 self._jitter([x / 255 for x in self.metadata.thing_colors[c]])572 for c in category_ids573 ]574 names = self.metadata.get("thing_classes", None)575 labels = _create_text_labels(576 category_ids,577 scores=None,578 class_names=names,579 is_crowd=[x.get("iscrowd", 0) for x in annos],580 )581 self.overlay_instances(582 labels=labels, boxes=boxes, masks=masks, keypoints=keypts, assigned_colors=colors583 )584 585 sem_seg = dic.get("sem_seg", None)586 if sem_seg is None and "sem_seg_file_name" in dic:587 with PathManager.open(dic["sem_seg_file_name"], "rb") as f:588 sem_seg = Image.open(f)589 sem_seg = np.asarray(sem_seg, dtype="uint8")590 if sem_seg is not None:591 self.draw_sem_seg(sem_seg, area_threshold=0, alpha=0.5)592 593 pan_seg = dic.get("pan_seg", None)594 if pan_seg is None and "pan_seg_file_name" in dic:595 with PathManager.open(dic["pan_seg_file_name"], "rb") as f:596 pan_seg = Image.open(f)597 pan_seg = np.asarray(pan_seg)598 from panopticapi.utils import rgb2id599 600 pan_seg = rgb2id(pan_seg)601 if pan_seg is not None:602 segments_info = dic["segments_info"]603 pan_seg = torch.tensor(pan_seg)604 self.draw_panoptic_seg(pan_seg, segments_info, area_threshold=0, alpha=0.5)605 return self.output606 607 def overlay_instances(608 self,609 *,610 boxes=None,611 labels=None,612 masks=None,613 keypoints=None,614 assigned_colors=None,615 alpha=0.5,616 ):617 """618 Args:619 boxes (Boxes, RotatedBoxes or ndarray): either a :class:`Boxes`,620 or an Nx4 numpy array of XYXY_ABS format for the N objects in a single image,621 or a :class:`RotatedBoxes`,622 or an Nx5 numpy array of (x_center, y_center, width, height, angle_degrees) format623 for the N objects in a single image,624 labels (list[str]): the text to be displayed for each instance.625 masks (masks-like object): Supported types are:626 627 * :class:`detectron2.structures.PolygonMasks`,628 :class:`detectron2.structures.BitMasks`.629 * list[list[ndarray]]: contains the segmentation masks for all objects in one image.630 The first level of the list corresponds to individual instances. The second631 level to all the polygon that compose the instance, and the third level632 to the polygon coordinates. The third level should have the format of633 [x0, y0, x1, y1, ..., xn, yn] (n >= 3).634 * list[ndarray]: each ndarray is a binary mask of shape (H, W).635 * list[dict]: each dict is a COCO-style RLE.636 keypoints (Keypoint or array like): an array-like object of shape (N, K, 3),637 where the N is the number of instances and K is the number of keypoints.638 The last dimension corresponds to (x, y, visibility or score).639 assigned_colors (list[matplotlib.colors]): a list of colors, where each color640 corresponds to each mask or box in the image. Refer to 'matplotlib.colors'641 for full list of formats that the colors are accepted in.642 Returns:643 output (VisImage): image object with visualizations.644 """645 num_instances = 0646 if boxes is not None:647 boxes = self._convert_boxes(boxes)648 num_instances = len(boxes)649 if masks is not None:650 masks = self._convert_masks(masks)651 if num_instances:652 assert len(masks) == num_instances653 else:654 num_instances = len(masks)655 if keypoints is not None:656 if num_instances:657 assert len(keypoints) == num_instances658 else:659 num_instances = len(keypoints)660 keypoints = self._convert_keypoints(keypoints)661 if labels is not None:662 assert len(labels) == num_instances663 if assigned_colors is None:664 assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)]665 if num_instances == 0:666 return self.output667 if boxes is not None and boxes.shape[1] == 5:668 return self.overlay_rotated_instances(669 boxes=boxes, labels=labels, assigned_colors=assigned_colors670 )671 672 # Display in largest to smallest order to reduce occlusion.673 areas = None674 if boxes is not None:675 areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1)676 elif masks is not None:677 areas = np.asarray([x.area() for x in masks])678 679 if areas is not None:680 sorted_idxs = np.argsort(-areas).tolist()681 # Re-order overlapped instances in descending order.682 boxes = boxes[sorted_idxs] if boxes is not None else None683 labels = [labels[k] for k in sorted_idxs] if labels is not None else None684 masks = [masks[idx] for idx in sorted_idxs] if masks is not None else None685 assigned_colors = [assigned_colors[idx] for idx in sorted_idxs]686 keypoints = keypoints[sorted_idxs] if keypoints is not None else None687 688 for i in range(num_instances):689 color = assigned_colors[i]690 if boxes is not None:691 self.draw_box(boxes[i], edge_color=color)692 693 if masks is not None:694 for segment in masks[i].polygons:695 self.draw_polygon(segment.reshape(-1, 2), color, alpha=alpha)696 697 if labels is not None:698 # first get a box699 if boxes is not None:700 x0, y0, x1, y1 = boxes[i]701 text_pos = (x0, y0) # if drawing boxes, put text on the box corner.702 horiz_align = "left"703 elif masks is not None:704 # skip small mask without polygon705 if len(masks[i].polygons) == 0:706 continue707 708 x0, y0, x1, y1 = masks[i].bbox()709 710 # draw text in the center (defined by median) when box is not drawn711 # median is less sensitive to outliers.712 text_pos = np.median(masks[i].mask.nonzero(), axis=1)[::-1]713 horiz_align = "center"714 else:715 continue # drawing the box confidence for keypoints isn't very useful.716 # for small objects, draw text at the side to avoid occlusion717 instance_area = (y1 - y0) * (x1 - x0)718 if (719 instance_area < _SMALL_OBJECT_AREA_THRESH * self.output.scale720 or y1 - y0 < 40 * self.output.scale721 ):722 if y1 >= self.output.height - 5:723 text_pos = (x1, y0)724 else:725 text_pos = (x0, y1)726 727 height_ratio = (y1 - y0) / np.sqrt(self.output.height * self.output.width)728 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)729 font_size = (730 np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2)731 * 0.5732 * self._default_font_size733 )734 self.draw_text(735 labels[i],736 text_pos,737 color=lighter_color,738 horizontal_alignment=horiz_align,739 font_size=font_size,740 )741 742 # draw keypoints743 if keypoints is not None:744 for keypoints_per_instance in keypoints:745 self.draw_and_connect_keypoints(keypoints_per_instance)746 747 return self.output748 749 def overlay_rotated_instances(self, boxes=None, labels=None, assigned_colors=None):750 """751 Args:752 boxes (ndarray): an Nx5 numpy array of753 (x_center, y_center, width, height, angle_degrees) format754 for the N objects in a single image.755 labels (list[str]): the text to be displayed for each instance.756 assigned_colors (list[matplotlib.colors]): a list of colors, where each color757 corresponds to each mask or box in the image. Refer to 'matplotlib.colors'758 for full list of formats that the colors are accepted in.759 760 Returns:761 output (VisImage): image object with visualizations.762 """763 num_instances = len(boxes)764 765 if assigned_colors is None:766 assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)]767 if num_instances == 0:768 return self.output769 770 # Display in largest to smallest order to reduce occlusion.771 if boxes is not None:772 areas = boxes[:, 2] * boxes[:, 3]773 774 sorted_idxs = np.argsort(-areas).tolist()775 # Re-order overlapped instances in descending order.776 boxes = boxes[sorted_idxs]777 labels = [labels[k] for k in sorted_idxs] if labels is not None else None778 colors = [assigned_colors[idx] for idx in sorted_idxs]779 780 for i in range(num_instances):781 self.draw_rotated_box_with_label(782 boxes[i], edge_color=colors[i], label=labels[i] if labels is not None else None783 )784 785 return self.output786 787 def draw_and_connect_keypoints(self, keypoints):788 """789 Draws keypoints of an instance and follows the rules for keypoint connections790 to draw lines between appropriate keypoints. This follows color heuristics for791 line color.792 793 Args:794 keypoints (Tensor): a tensor of shape (K, 3), where K is the number of keypoints795 and the last dimension corresponds to (x, y, probability).796 797 Returns:798 output (VisImage): image object with visualizations.799 """800 visible = {}801 keypoint_names = self.metadata.get("keypoint_names")802 for idx, keypoint in enumerate(keypoints):803 804 # draw keypoint805 x, y, prob = keypoint806 if prob > self.keypoint_threshold:807 self.draw_circle((x, y), color=_RED)808 if keypoint_names:809 keypoint_name = keypoint_names[idx]810 visible[keypoint_name] = (x, y)811 812 if self.metadata.get("keypoint_connection_rules"):813 for kp0, kp1, color in self.metadata.keypoint_connection_rules:814 if kp0 in visible and kp1 in visible:815 x0, y0 = visible[kp0]816 x1, y1 = visible[kp1]817 color = tuple(x / 255.0 for x in color)818 self.draw_line([x0, x1], [y0, y1], color=color)819 820 # draw lines from nose to mid-shoulder and mid-shoulder to mid-hip821 # Note that this strategy is specific to person keypoints.822 # For other keypoints, it should just do nothing823 try:824 ls_x, ls_y = visible["left_shoulder"]825 rs_x, rs_y = visible["right_shoulder"]826 mid_shoulder_x, mid_shoulder_y = (ls_x + rs_x) / 2, (ls_y + rs_y) / 2827 except KeyError:828 pass829 else:830 # draw line from nose to mid-shoulder831 nose_x, nose_y = visible.get("nose", (None, None))832 if nose_x is not None:833 self.draw_line([nose_x, mid_shoulder_x], [nose_y, mid_shoulder_y], color=_RED)834 835 try:836 # draw line from mid-shoulder to mid-hip837 lh_x, lh_y = visible["left_hip"]838 rh_x, rh_y = visible["right_hip"]839 except KeyError:840 pass841 else:842 mid_hip_x, mid_hip_y = (lh_x + rh_x) / 2, (lh_y + rh_y) / 2843 self.draw_line([mid_hip_x, mid_shoulder_x], [mid_hip_y, mid_shoulder_y], color=_RED)844 return self.output845 846 """847 Primitive drawing functions:848 """849 850 def draw_text(851 self,852 text,853 position,854 *,855 font_size=None,856 color="g",857 horizontal_alignment="center",858 rotation=0,859 ):860 """861 Args:862 text (str): class label863 position (tuple): a tuple of the x and y coordinates to place text on image.864 font_size (int, optional): font of the text. If not provided, a font size865 proportional to the image width is calculated and used.866 color: color of the text. Refer to `matplotlib.colors` for full list867 of formats that are accepted.868 horizontal_alignment (str): see `matplotlib.text.Text`869 rotation: rotation angle in degrees CCW870 871 Returns:872 output (VisImage): image object with text drawn.873 """874 if not font_size:875 font_size = self._default_font_size876 877 # since the text background is dark, we don't want the text to be dark878 color = np.maximum(list(mplc.to_rgb(color)), 0.2)879 color[np.argmax(color)] = max(0.8, np.max(color))880 881 x, y = position882 self.output.ax.text(883 x,884 y,885 text,886 size=font_size * self.output.scale,887 family="sans-serif",888 bbox={"facecolor": "black", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"},889 verticalalignment="top",890 horizontalalignment=horizontal_alignment,891 color=color,892 zorder=10,893 rotation=rotation,894 )895 return self.output896 897 def draw_box(self, box_coord, alpha=0.5, edge_color="g", line_style="-"):898 """899 Args:900 box_coord (tuple): a tuple containing x0, y0, x1, y1 coordinates, where x0 and y0901 are the coordinates of the image's top left corner. x1 and y1 are the902 coordinates of the image's bottom right corner.903 alpha (float): blending efficient. Smaller values lead to more transparent masks.904 edge_color: color of the outline of the box. Refer to `matplotlib.colors`905 for full list of formats that are accepted.906 line_style (string): the string to use to create the outline of the boxes.907 908 Returns:909 output (VisImage): image object with box drawn.910 """911 x0, y0, x1, y1 = box_coord912 width = x1 - x0913 height = y1 - y0914 915 linewidth = max(self._default_font_size / 4, 1)916 917 self.output.ax.add_patch(918 mpl.patches.Rectangle(919 (x0, y0),920 width,921 height,922 fill=False,923 edgecolor=edge_color,924 linewidth=linewidth * self.output.scale,925 alpha=alpha,926 linestyle=line_style,927 )928 )929 return self.output930 931 def draw_rotated_box_with_label(932 self, rotated_box, alpha=0.5, edge_color="g", line_style="-", label=None933 ):934 """935 Draw a rotated box with label on its top-left corner.936 937 Args:938 rotated_box (tuple): a tuple containing (cnt_x, cnt_y, w, h, angle),939 where cnt_x and cnt_y are the center coordinates of the box.940 w and h are the width and height of the box. angle represents how941 many degrees the box is rotated CCW with regard to the 0-degree box.942 alpha (float): blending efficient. Smaller values lead to more transparent masks.943 edge_color: color of the outline of the box. Refer to `matplotlib.colors`944 for full list of formats that are accepted.945 line_style (string): the string to use to create the outline of the boxes.946 label (string): label for rotated box. It will not be rendered when set to None.947 948 Returns:949 output (VisImage): image object with box drawn.950 """951 cnt_x, cnt_y, w, h, angle = rotated_box952 area = w * h953 # use thinner lines when the box is small954 linewidth = self._default_font_size / (955 6 if area < _SMALL_OBJECT_AREA_THRESH * self.output.scale else 3956 )957 958 theta = angle * math.pi / 180.0959 c = math.cos(theta)960 s = math.sin(theta)961 rect = [(-w / 2, h / 2), (-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2)]962 # x: left->right ; y: top->down963 rotated_rect = [(s * yy + c * xx + cnt_x, c * yy - s * xx + cnt_y) for (xx, yy) in rect]964 for k in range(4):965 j = (k + 1) % 4966 self.draw_line(967 [rotated_rect[k][0], rotated_rect[j][0]],968 [rotated_rect[k][1], rotated_rect[j][1]],969 color=edge_color,970 linestyle="--" if k == 1 else line_style,971 linewidth=linewidth,972 )973 974 if label is not None:975 text_pos = rotated_rect[1] # topleft corner976 977 height_ratio = h / np.sqrt(self.output.height * self.output.width)978 label_color = self._change_color_brightness(edge_color, brightness_factor=0.7)979 font_size = (980 np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) * 0.5 * self._default_font_size981 )982 self.draw_text(label, text_pos, color=label_color, font_size=font_size, rotation=angle)983 984 return self.output985 986 def draw_circle(self, circle_coord, color, radius=3):987 """988 Args:989 circle_coord (list(int) or tuple(int)): contains the x and y coordinates990 of the center of the circle.991 color: color of the polygon. Refer to `matplotlib.colors` for a full list of992 formats that are accepted.993 radius (int): radius of the circle.994 995 Returns:996 output (VisImage): image object with box drawn.997 """998 x, y = circle_coord999 self.output.ax.add_patch(1000 mpl.patches.Circle(circle_coord, radius=radius, fill=True, color=color)1001 )1002 return self.output1003 1004 def draw_line(self, x_data, y_data, color, linestyle="-", linewidth=None):1005 """1006 Args:1007 x_data (list[int]): a list containing x values of all the points being drawn.1008 Length of list should match the length of y_data.1009 y_data (list[int]): a list containing y values of all the points being drawn.1010 Length of list should match the length of x_data.1011 color: color of the line. Refer to `matplotlib.colors` for a full list of1012 formats that are accepted.1013 linestyle: style of the line. Refer to `matplotlib.lines.Line2D`1014 for a full list of formats that are accepted.1015 linewidth (float or None): width of the line. When it's None,1016 a default value will be computed and used.1017 1018 Returns:1019 output (VisImage): image object with line drawn.1020 """1021 if linewidth is None:1022 linewidth = self._default_font_size / 31023 linewidth = max(linewidth, 1)1024 self.output.ax.add_line(1025 mpl.lines.Line2D(1026 x_data,1027 y_data,1028 linewidth=linewidth * self.output.scale,1029 color=color,1030 linestyle=linestyle,1031 )1032 )1033 return self.output1034 1035 def draw_binary_mask(1036 self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=101037 ):1038 """1039 Args:1040 binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and1041 W is the image width. Each value in the array is either a 0 or 1 value of uint81042 type.1043 color: color of the mask. Refer to `matplotlib.colors` for a full list of1044 formats that are accepted. If None, will pick a random color.1045 edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a1046 full list of formats that are accepted.1047 text (str): if None, will be drawn on the object1048 alpha (float): blending efficient. Smaller values lead to more transparent masks.1049 area_threshold (float): a connected component smaller than this area will not be shown.1050 1051 Returns:1052 output (VisImage): image object with mask drawn.1053 """1054 if color is None:1055 color = random_color(rgb=True, maximum=1)1056 color = mplc.to_rgb(color)1057 1058 has_valid_segment = False1059 binary_mask = binary_mask.astype("uint8") # opencv needs uint81060 mask = GenericMask(binary_mask, self.output.height, self.output.width)1061 shape2d = (binary_mask.shape[0], binary_mask.shape[1])1062 1063 if not mask.has_holes:1064 # draw polygons for regular masks1065 for segment in mask.polygons:1066 area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1]))1067 if area < (area_threshold or 0):1068 continue1069 has_valid_segment = True1070 segment = segment.reshape(-1, 2)1071 self.draw_polygon(segment, color=color, edge_color=edge_color, alpha=alpha)1072 else:1073 # TODO: Use Path/PathPatch to draw vector graphics:1074 # https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon1075 rgba = np.zeros(shape2d + (4,), dtype="float32")1076 rgba[:, :, :3] = color1077 rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha1078 has_valid_segment = True1079 self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0))1080 1081 if text is not None and has_valid_segment:1082 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)1083 self._draw_text_in_mask(binary_mask, text, lighter_color)1084 return self.output1085 1086 def draw_soft_mask(self, soft_mask, color=None, *, text=None, alpha=0.5):1087 """1088 Args:1089 soft_mask (ndarray): float array of shape (H, W), each value in [0, 1].1090 color: color of the mask. Refer to `matplotlib.colors` for a full list of1091 formats that are accepted. If None, will pick a random color.1092 text (str): if None, will be drawn on the object1093 alpha (float): blending efficient. Smaller values lead to more transparent masks.1094 1095 Returns:1096 output (VisImage): image object with mask drawn.1097 """1098 if color is None:1099 color = random_color(rgb=True, maximum=1)1100 color = mplc.to_rgb(color)1101 1102 shape2d = (soft_mask.shape[0], soft_mask.shape[1])1103 rgba = np.zeros(shape2d + (4,), dtype="float32")1104 rgba[:, :, :3] = color1105 rgba[:, :, 3] = soft_mask * alpha1106 self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0))1107 1108 if text is not None:1109 lighter_color = self._change_color_brightness(color, brightness_factor=0.7)1110 binary_mask = (soft_mask > 0.5).astype("uint8")1111 self._draw_text_in_mask(binary_mask, text, lighter_color)1112 return self.output1113 1114 def draw_polygon(self, segment, color, edge_color=None, alpha=0.5):1115 """1116 Args:1117 segment: numpy array of shape Nx2, containing all the points in the polygon.1118 color: color of the polygon. Refer to `matplotlib.colors` for a full list of1119 formats that are accepted.1120 edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a1121 full list of formats that are accepted. If not provided, a darker shade1122 of the polygon color will be used instead.1123 alpha (float): blending efficient. Smaller values lead to more transparent masks.1124 1125 Returns:1126 output (VisImage): image object with polygon drawn.1127 """1128 if edge_color is None:1129 # make edge color darker than the polygon color1130 if alpha > 0.8:1131 edge_color = self._change_color_brightness(color, brightness_factor=-0.7)1132 else:1133 edge_color = color1134 edge_color = mplc.to_rgb(edge_color) + (1,)1135 1136 polygon = mpl.patches.Polygon(1137 segment,1138 fill=True,1139 facecolor=mplc.to_rgb(color) + (alpha,),1140 edgecolor=edge_color,1141 linewidth=max(self._default_font_size // 15 * self.output.scale, 1),1142 )1143 self.output.ax.add_patch(polygon)1144 return self.output1145 1146 """1147 Internal methods:1148 """1149 1150 def _jitter(self, color):1151 """1152 Randomly modifies given color to produce a slightly different color than the color given.1153 1154 Args:1155 color (tuple[double]): a tuple of 3 elements, containing the RGB values of the color1156 picked. The values in the list are in the [0.0, 1.0] range.1157 1158 Returns:1159 jittered_color (tuple[double]): a tuple of 3 elements, containing the RGB values of the1160 color after being jittered. The values in the list are in the [0.0, 1.0] range.1161 """1162 color = mplc.to_rgb(color)1163 vec = np.random.rand(3)1164 # better to do it in another color space1165 vec = vec / np.linalg.norm(vec) * 0.51166 res = np.clip(vec + color, 0, 1)1167 return tuple(res)1168 1169 def _create_grayscale_image(self, mask=None):1170 """1171 Create a grayscale version of the original image.1172 The colors in masked area, if given, will be kept.1173 """1174 img_bw = self.img.astype("f4").mean(axis=2)1175 img_bw = np.stack([img_bw] * 3, axis=2)1176 if mask is not None:1177 img_bw[mask] = self.img[mask]1178 return img_bw1179 1180 def _change_color_brightness(self, color, brightness_factor):1181 """1182 Depending on the brightness_factor, gives a lighter or darker color i.e. a color with1183 less or more saturation than the original color.1184 1185 Args:1186 color: color of the polygon. Refer to `matplotlib.colors` for a full list of1187 formats that are accepted.1188 brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of1189 0 will correspond to no change, a factor in [-1.0, 0) range will result in1190 a darker color and a factor in (0, 1.0] range will result in a lighter color.1191 1192 Returns:1193 modified_color (tuple[double]): a tuple containing the RGB values of the1194 modified color. Each value in the tuple is in the [0.0, 1.0] range.1195 """1196 assert brightness_factor >= -1.0 and brightness_factor <= 1.01197 color = mplc.to_rgb(color)1198 polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color))1199 modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1])1200 modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness