Spaces:
Build error
Build error
# Edit by Yao Lu | |
# | |
# Copyright (c) Facebook, Inc. and its affiliates. | |
import colorsys | |
import logging | |
import math | |
import numpy as np | |
from enum import Enum, unique | |
import cv2 | |
import matplotlib as mpl | |
import matplotlib.colors as mplc | |
import matplotlib.figure as mplfigure | |
import pycocotools.mask as mask_util | |
import torch | |
from matplotlib.backends.backend_agg import FigureCanvasAgg | |
from PIL import Image | |
from detectron2.data import MetadataCatalog | |
from detectron2.structures import BitMasks, Boxes, BoxMode, Keypoints, PolygonMasks, RotatedBoxes | |
from detectron2.utils.file_io import PathManager | |
from .colormap import random_color | |
from shapely.geometry import * | |
import pickle | |
import matplotlib.font_manager as mfm | |
logger = logging.getLogger(__name__) | |
__all__ = ["ColorMode", "VisImage", "Visualizer"] | |
_SMALL_OBJECT_AREA_THRESH = 1000 | |
_LARGE_MASK_AREA_THRESH = 120000 | |
_OFF_WHITE = (1.0, 1.0, 240.0 / 255) | |
_BLACK = (0, 0, 0) | |
_RED = (1.0, 0, 0) | |
_KEYPOINT_THRESHOLD = 0.05 | |
def py_cpu_pnms(dets, scores, thresh): | |
pts = dets | |
# for i in xrange(dets.shape[0]): | |
# pts.append([[int(bbox[i, 0]) + info_bbox[i, j], int(bbox[i, 1]) + info_bbox[i, j+1]] for j in xrange(0,28,2)]) | |
scores = np.array(scores) | |
order = scores.argsort()[::-1] | |
areas = np.zeros(scores.shape) | |
order = scores.argsort()[::-1] | |
inter_areas = np.zeros((scores.shape[0], scores.shape[0])) | |
for il in range(len(pts)): | |
poly = Polygon(pts[il]).buffer(0.001) | |
areas[il] = poly.area | |
for jl in range(il, len(pts)): | |
polyj = Polygon(pts[jl].tolist()).buffer(0.001) | |
inS = poly.intersection(polyj) | |
try: | |
inter_areas[il][jl] = inS.area | |
except: | |
import pdb;pdb.set_trace() | |
inter_areas[jl][il] = inS.area | |
keep = [] | |
while order.size > 0: | |
i = order[0] | |
keep.append(i) | |
ovr = inter_areas[i][order[1:]] / ((areas[i]) + areas[order[1:]] - inter_areas[i][order[1:]]) | |
inds = np.where(ovr <= thresh)[0] | |
order = order[inds + 1] | |
return keep | |
class ColorMode(Enum): | |
""" | |
Enum of different color modes to use for instance visualizations. | |
""" | |
IMAGE = 0 | |
""" | |
Picks a random color for every instance and overlay segmentations with low opacity. | |
""" | |
SEGMENTATION = 1 | |
""" | |
Let instances of the same category have similar colors | |
(from metadata.thing_colors), and overlay them with | |
high opacity. This provides more attention on the quality of segmentation. | |
""" | |
IMAGE_BW = 2 | |
""" | |
Same as IMAGE, but convert all areas without masks to gray-scale. | |
Only available for drawing per-instance mask predictions. | |
""" | |
class GenericMask: | |
""" | |
Attribute: | |
polygons (list[ndarray]): list[ndarray]: polygons for this mask. | |
Each ndarray has format [x, y, x, y, ...] | |
mask (ndarray): a binary mask | |
""" | |
def __init__(self, mask_or_polygons, height, width): | |
self._mask = self._polygons = self._has_holes = None | |
self.height = height | |
self.width = width | |
m = mask_or_polygons | |
if isinstance(m, dict): | |
# RLEs | |
assert "counts" in m and "size" in m | |
if isinstance(m["counts"], list): # uncompressed RLEs | |
h, w = m["size"] | |
assert h == height and w == width | |
m = mask_util.frPyObjects(m, h, w) | |
self._mask = mask_util.decode(m)[:, :] | |
return | |
if isinstance(m, list): # list[ndarray] | |
self._polygons = [np.asarray(x).reshape(-1) for x in m] | |
return | |
if isinstance(m, np.ndarray): # assumed to be a binary mask | |
assert m.shape[1] != 2, m.shape | |
assert m.shape == (height, width), m.shape | |
self._mask = m.astype("uint8") | |
return | |
raise ValueError("GenericMask cannot handle object {} of type '{}'".format(m, type(m))) | |
def mask(self): | |
if self._mask is None: | |
self._mask = self.polygons_to_mask(self._polygons) | |
return self._mask | |
def polygons(self): | |
if self._polygons is None: | |
self._polygons, self._has_holes = self.mask_to_polygons(self._mask) | |
return self._polygons | |
def has_holes(self): | |
if self._has_holes is None: | |
if self._mask is not None: | |
self._polygons, self._has_holes = self.mask_to_polygons(self._mask) | |
else: | |
self._has_holes = False # if original format is polygon, does not have holes | |
return self._has_holes | |
def mask_to_polygons(self, mask): | |
# cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level | |
# hierarchy. External contours (boundary) of the object are placed in hierarchy-1. | |
# Internal contours (holes) are placed in hierarchy-2. | |
# cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours. | |
mask = np.ascontiguousarray(mask) # some versions of cv2 does not support incontiguous arr | |
#res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE) | |
res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) | |
hierarchy = res[-1] | |
if hierarchy is None: # empty mask | |
return [], False | |
has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0 | |
res = res[-2] | |
res = [x.flatten() for x in res] | |
# These coordinates from OpenCV are integers in range [0, W-1 or H-1]. | |
# We add 0.5 to turn them into real-value coordinate space. A better solution | |
# would be to first +0.5 and then dilate the returned polygon by 0.5. | |
res = [x + 0.5 for x in res if len(x) >= 6] | |
return res, has_holes | |
def polygons_to_mask(self, polygons): | |
rle = mask_util.frPyObjects(polygons, self.height, self.width) | |
rle = mask_util.merge(rle) | |
return mask_util.decode(rle)[:, :] | |
def area(self): | |
return self.mask.sum() | |
def bbox(self): | |
p = mask_util.frPyObjects(self.polygons, self.height, self.width) | |
p = mask_util.merge(p) | |
bbox = mask_util.toBbox(p) | |
bbox[2] += bbox[0] | |
bbox[3] += bbox[1] | |
return bbox | |
class _PanopticPrediction: | |
def __init__(self, panoptic_seg, segments_info, metadata=None): | |
if segments_info is None: | |
assert metadata is not None | |
# If "segments_info" is None, we assume "panoptic_img" is a | |
# H*W int32 image storing the panoptic_id in the format of | |
# category_id * label_divisor + instance_id. We reserve -1 for | |
# VOID label. | |
label_divisor = metadata.label_divisor | |
segments_info = [] | |
for panoptic_label in np.unique(panoptic_seg.numpy()): | |
if panoptic_label == -1: | |
# VOID region. | |
continue | |
pred_class = panoptic_label // label_divisor | |
isthing = pred_class in metadata.thing_dataset_id_to_contiguous_id.values() | |
segments_info.append( | |
{ | |
"id": int(panoptic_label), | |
"category_id": int(pred_class), | |
"isthing": bool(isthing), | |
} | |
) | |
del metadata | |
self._seg = panoptic_seg | |
self._sinfo = {s["id"]: s for s in segments_info} # seg id -> seg info | |
segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True) | |
areas = areas.numpy() | |
sorted_idxs = np.argsort(-areas) | |
self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs] | |
self._seg_ids = self._seg_ids.tolist() | |
for sid, area in zip(self._seg_ids, self._seg_areas): | |
if sid in self._sinfo: | |
self._sinfo[sid]["area"] = float(area) | |
def non_empty_mask(self): | |
""" | |
Returns: | |
(H, W) array, a mask for all pixels that have a prediction | |
""" | |
empty_ids = [] | |
for id in self._seg_ids: | |
if id not in self._sinfo: | |
empty_ids.append(id) | |
if len(empty_ids) == 0: | |
return np.zeros(self._seg.shape, dtype=np.uint8) | |
assert ( | |
len(empty_ids) == 1 | |
), ">1 ids corresponds to no labels. This is currently not supported" | |
return (self._seg != empty_ids[0]).numpy().astype(np.bool) | |
def semantic_masks(self): | |
for sid in self._seg_ids: | |
sinfo = self._sinfo.get(sid) | |
if sinfo is None or sinfo["isthing"]: | |
# Some pixels (e.g. id 0 in PanopticFPN) have no instance or semantic predictions. | |
continue | |
yield (self._seg == sid).numpy().astype(np.bool), sinfo | |
def instance_masks(self): | |
for sid in self._seg_ids: | |
sinfo = self._sinfo.get(sid) | |
if sinfo is None or not sinfo["isthing"]: | |
continue | |
mask = (self._seg == sid).numpy().astype(np.bool) | |
if mask.sum() > 0: | |
yield mask, sinfo | |
def _create_text_labels(classes, scores, class_names): | |
""" | |
Args: | |
classes (list[int] or None): | |
scores (list[float] or None): | |
class_names (list[str] or None): | |
Returns: | |
list[str] or None | |
""" | |
labels = None | |
if classes is not None and class_names is not None and len(class_names) > 0: | |
labels = [class_names[i] for i in classes] | |
if scores is not None: | |
if labels is None: | |
labels = ["{:.0f}%".format(s * 100) for s in scores] | |
else: | |
# labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)] | |
#luyao | |
labels = ["{}.{:.0f}".format(l, s * 100) for l, s in zip(labels, scores)] | |
return labels | |
class VisImage: | |
def __init__(self, img, scale=1.0): | |
""" | |
Args: | |
img (ndarray): an RGB image of shape (H, W, 3). | |
scale (float): scale the input image | |
""" | |
self.img = img | |
self.scale = scale | |
self.width, self.height = img.shape[1], img.shape[0] | |
self._setup_figure(img) | |
def _setup_figure(self, img): | |
""" | |
Args: | |
Same as in :meth:`__init__()`. | |
Returns: | |
fig (matplotlib.pyplot.figure): top level container for all the image plot elements. | |
ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system. | |
""" | |
fig = mplfigure.Figure(frameon=False) | |
self.dpi = fig.get_dpi() | |
# add a small 1e-2 to avoid precision lost due to matplotlib's truncation | |
# (https://github.com/matplotlib/matplotlib/issues/15363) | |
fig.set_size_inches( | |
(self.width * self.scale + 1e-2) / self.dpi, | |
(self.height * self.scale + 1e-2) / self.dpi, | |
) | |
self.canvas = FigureCanvasAgg(fig) | |
# self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig) | |
ax = fig.add_axes([0.0, 0.0, 1.0, 1.0]) | |
ax.axis("off") | |
# Need to imshow this first so that other patches can be drawn on top | |
ax.imshow(img, extent=(0, self.width, self.height, 0), interpolation="nearest") | |
self.fig = fig | |
self.ax = ax | |
def save(self, filepath): | |
""" | |
Args: | |
filepath (str): a string that contains the absolute path, including the file name, where | |
the visualized image will be saved. | |
""" | |
self.fig.savefig(filepath) | |
# self.fig.savefig(filepath[:-4]+'.svg', format='svg') | |
def get_image(self): | |
""" | |
Returns: | |
ndarray: | |
the visualized image of shape (H, W, 3) (RGB) in uint8 type. | |
The shape is scaled w.r.t the input image using the given `scale` argument. | |
""" | |
canvas = self.canvas | |
s, (width, height) = canvas.print_to_buffer() | |
# buf = io.BytesIO() # works for cairo backend | |
# canvas.print_rgba(buf) | |
# width, height = self.width, self.height | |
# s = buf.getvalue() | |
buffer = np.frombuffer(s, dtype="uint8") | |
img_rgba = buffer.reshape(height, width, 4) | |
rgb, alpha = np.split(img_rgba, [3], axis=2) | |
return rgb.astype("uint8") | |
class Visualizer: | |
""" | |
Visualizer that draws data about detection/segmentation on images. | |
It contains methods like `draw_{text,box,circle,line,binary_mask,polygon}` | |
that draw primitive objects to images, as well as high-level wrappers like | |
`draw_{instance_predictions,sem_seg,panoptic_seg_predictions,dataset_dict}` | |
that draw composite data in some pre-defined style. | |
Note that the exact visualization style for the high-level wrappers are subject to change. | |
Style such as color, opacity, label contents, visibility of labels, or even the visibility | |
of objects themselves (e.g. when the object is too small) may change according | |
to different heuristics, as long as the results still look visually reasonable. | |
To obtain a consistent style, implement custom drawing functions with the primitive | |
methods instead. | |
This visualizer focuses on high rendering quality rather than performance. It is not | |
designed to be used for real-time applications. | |
""" | |
# TODO implement a fast, rasterized version using OpenCV | |
def __init__(self, img_rgb, metadata=None, scale=1.0, instance_mode=ColorMode.IMAGE): | |
""" | |
Args: | |
img_rgb: a numpy array of shape (H, W, C), where H and W correspond to | |
the height and width of the image respectively. C is the number of | |
color channels. The image is required to be in RGB format since that | |
is a requirement of the Matplotlib library. The image is also expected | |
to be in the range [0, 255]. | |
metadata (Metadata): image metadata. | |
instance_mode (ColorMode): defines one of the pre-defined style for drawing | |
instances on an image. | |
""" | |
self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8) | |
if metadata is None: | |
metadata = MetadataCatalog.get("__nonexist__") | |
self.metadata = metadata | |
self.output = VisImage(self.img, scale=scale) | |
self.cpu_device = torch.device("cpu") | |
# too small texts are useless, therefore clamp to 9 | |
self._default_font_size = max( | |
np.sqrt(self.output.height * self.output.width) // 90, 10 // scale | |
) | |
self._instance_mode = instance_mode | |
with open('chn_cls_list.txt', 'rb') as fp: | |
self.CTLABELS = pickle.load(fp) | |
def draw_instance_predictions(self, predictions, path): | |
""" | |
Draw instance-level prediction results on an image. | |
Args: | |
predictions (Instances): the output of an instance detection/segmentation | |
model. Following fields will be used to draw: | |
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle"). | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None | |
scores = predictions.scores if predictions.has("scores") else None | |
classes = predictions.pred_classes if predictions.has("pred_classes") else None | |
labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None)) | |
#luyao# | |
# labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None)) | |
keypoints = predictions.pred_keypoints if predictions.has("pred_keypoints") else None | |
rec = predictions.pred_rec if predictions.has("pred_rec") else None | |
rec_score = predictions.pred_rec_score if predictions.has("pred_rec_score") else None | |
#luyao# | |
if predictions.has("pred_masks"): | |
masks = np.asarray(predictions.pred_masks) | |
masks = [GenericMask(x, self.output.height, self.output.width) for x in masks] | |
else: | |
masks = None | |
# masks = None | |
if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"): | |
colors = [ | |
self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in classes | |
] | |
#luyao# | |
alpha = 0.8 | |
else: | |
colors = None | |
alpha = 0.77 | |
if self._instance_mode == ColorMode.IMAGE_BW: | |
self.output.img = self._create_grayscale_image( | |
(predictions.pred_masks.any(dim=0) > 0).numpy() | |
if predictions.has("pred_masks") | |
else None | |
) | |
alpha = 0.3 | |
self.overlay_instances( | |
rec=rec, | |
masks=masks, | |
boxes=boxes, | |
labels=labels, | |
keypoints=keypoints, | |
assigned_colors=colors, | |
alpha=alpha, | |
scores=scores, | |
path=path, | |
rec_score = rec_score | |
) | |
return self.output | |
def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8): | |
""" | |
Draw semantic segmentation predictions/labels. | |
Args: | |
sem_seg (Tensor or ndarray): the segmentation of shape (H, W). | |
Each value is the integer label of the pixel. | |
area_threshold (int): segments with less than `area_threshold` are not drawn. | |
alpha (float): the larger it is, the more opaque the segmentations are. | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
if isinstance(sem_seg, torch.Tensor): | |
sem_seg = sem_seg.numpy() | |
labels, areas = np.unique(sem_seg, return_counts=True) | |
sorted_idxs = np.argsort(-areas).tolist() | |
labels = labels[sorted_idxs] | |
for label in filter(lambda l: l < len(self.metadata.stuff_classes), labels): | |
try: | |
mask_color = [x / 255 for x in self.metadata.stuff_colors[label]] | |
except (AttributeError, IndexError): | |
mask_color = None | |
binary_mask = (sem_seg == label).astype(np.uint8) | |
text = self.metadata.stuff_classes[label] | |
self.draw_binary_mask( | |
binary_mask, | |
color=mask_color, | |
edge_color=_OFF_WHITE, | |
text=text, | |
alpha=alpha, | |
area_threshold=area_threshold, | |
) | |
return self.output | |
def draw_panoptic_seg_predictions( | |
self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7 | |
): | |
""" | |
Draw panoptic prediction results on an image. | |
Args: | |
panoptic_seg (Tensor): of shape (height, width) where the values are ids for each | |
segment. | |
segments_info (list[dict]): Describe each segment in `panoptic_seg`. | |
Each dict contains keys "id", "category_id", "isthing". | |
area_threshold (int): stuff segments with less than `area_threshold` are not drawn. | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata) | |
if self._instance_mode == ColorMode.IMAGE_BW: | |
self.output.img = self._create_grayscale_image(pred.non_empty_mask()) | |
# draw mask for all semantic segments first i.e. "stuff" | |
for mask, sinfo in pred.semantic_masks(): | |
category_idx = sinfo["category_id"] | |
try: | |
mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]] | |
except AttributeError: | |
mask_color = None | |
text = self.metadata.stuff_classes[category_idx] | |
self.draw_binary_mask( | |
mask, | |
color=mask_color, | |
edge_color=_OFF_WHITE, | |
text=text, | |
alpha=alpha, | |
area_threshold=area_threshold, | |
) | |
# draw mask for all instances second | |
all_instances = list(pred.instance_masks()) | |
if len(all_instances) == 0: | |
return self.output | |
masks, sinfo = list(zip(*all_instances)) | |
category_ids = [x["category_id"] for x in sinfo] | |
try: | |
scores = [x["score"] for x in sinfo] | |
except KeyError: | |
scores = None | |
labels = _create_text_labels(category_ids, scores, self.metadata.thing_classes) | |
try: | |
colors = [ | |
self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in category_ids | |
] | |
except AttributeError: | |
colors = None | |
self.overlay_instances(masks=masks, labels=labels, assigned_colors=colors, alpha=alpha) | |
return self.output | |
def draw_dataset_dict(self, dic): | |
""" | |
Draw annotations/segmentaions in Detectron2 Dataset format. | |
Args: | |
dic (dict): annotation/segmentation data of one image, in Detectron2 Dataset format. | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
annos = dic.get("annotations", None) | |
if annos: | |
if "segmentation" in annos[0]: | |
masks = [x["segmentation"] for x in annos] | |
else: | |
masks = None | |
if "keypoints" in annos[0]: | |
keypts = [x["keypoints"] for x in annos] | |
keypts = np.array(keypts).reshape(len(annos), -1, 3) | |
else: | |
keypts = None | |
boxes = [ | |
BoxMode.convert(x["bbox"], x["bbox_mode"], BoxMode.XYXY_ABS) | |
if len(x["bbox"]) == 4 | |
else x["bbox"] | |
for x in annos | |
] | |
labels = [x["category_id"] for x in annos] | |
colors = None | |
if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"): | |
colors = [ | |
self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in labels | |
] | |
names = self.metadata.get("thing_classes", None) | |
if names: | |
labels = [names[i] for i in labels] | |
labels = [ | |
"{}".format(i) + ("|crowd" if a.get("iscrowd", 0) else "") | |
for i, a in zip(labels, annos) | |
] | |
self.overlay_instances( | |
labels=labels, boxes=boxes, masks=masks, keypoints=keypts, assigned_colors=colors | |
) | |
sem_seg = dic.get("sem_seg", None) | |
if sem_seg is None and "sem_seg_file_name" in dic: | |
with PathManager.open(dic["sem_seg_file_name"], "rb") as f: | |
sem_seg = Image.open(f) | |
sem_seg = np.asarray(sem_seg, dtype="uint8") | |
if sem_seg is not None: | |
self.draw_sem_seg(sem_seg, area_threshold=0, alpha=0.5) | |
pan_seg = dic.get("pan_seg", None) | |
if pan_seg is None and "pan_seg_file_name" in dic: | |
assert "segments_info" in dic | |
with PathManager.open(dic["pan_seg_file_name"], "rb") as f: | |
pan_seg = Image.open(f) | |
pan_seg = np.asarray(pan_seg) | |
from panopticapi.utils import rgb2id | |
pan_seg = rgb2id(pan_seg) | |
segments_info = dic["segments_info"] | |
if pan_seg is not None: | |
pan_seg = torch.Tensor(pan_seg) | |
self.draw_panoptic_seg_predictions(pan_seg, segments_info, area_threshold=0, alpha=0.5) | |
return self.output | |
def overlay_instances( | |
self, | |
*, | |
rec=None, | |
boxes=None, | |
labels=None, | |
masks=None, | |
keypoints=None, | |
assigned_colors=None, | |
alpha=0.5, | |
scores, | |
path, | |
rec_score, | |
): | |
""" | |
Args: | |
boxes (Boxes, RotatedBoxes or ndarray): either a :class:`Boxes`, | |
or an Nx4 numpy array of XYXY_ABS format for the N objects in a single image, | |
or a :class:`RotatedBoxes`, | |
or an Nx5 numpy array of (x_center, y_center, width, height, angle_degrees) format | |
for the N objects in a single image, | |
labels (list[str]): the text to be displayed for each instance. | |
masks (masks-like object): Supported types are: | |
* :class:`detectron2.structures.PolygonMasks`, | |
:class:`detectron2.structures.BitMasks`. | |
* list[list[ndarray]]: contains the segmentation masks for all objects in one image. | |
The first level of the list corresponds to individual instances. The second | |
level to all the polygon that compose the instance, and the third level | |
to the polygon coordinates. The third level should have the format of | |
[x0, y0, x1, y1, ..., xn, yn] (n >= 3). | |
* list[ndarray]: each ndarray is a binary mask of shape (H, W). | |
* list[dict]: each dict is a COCO-style RLE. | |
keypoints (Keypoint or array like): an array-like object of shape (N, K, 3), | |
where the N is the number of instances and K is the number of keypoints. | |
The last dimension corresponds to (x, y, visibility or score). | |
assigned_colors (list[matplotlib.colors]): a list of colors, where each color | |
corresponds to each mask or box in the image. Refer to 'matplotlib.colors' | |
for full list of formats that the colors are accepted in. | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
rec = rec | |
def _ctc_decode_recognition(rec): | |
#CTLABELS = "_0123456789abcdefghijklmnopqrstuvwxyz" | |
# CTLABELS = [' ','!','"','#','$','%','&','\'','(',')','*','+',',','-','.','/','0','1','2','3','4','5','6','7','8','9',':',';','<','=','>','?','@','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z','[','\\',']','^','_','`','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z','{','|','}','~'] | |
# ctc decoding | |
s = '' | |
for c in rec: | |
c = int(c) | |
if c < 5461: | |
s += str(chr(self.CTLABELS[c])) | |
elif c == 5462: | |
s += u'' | |
return s | |
num_instances = None | |
if boxes is not None: | |
boxes = self._convert_boxes(boxes) | |
num_instances = len(boxes) | |
if masks is not None: | |
masks = self._convert_masks(masks) | |
if num_instances: | |
assert len(masks) == num_instances | |
else: | |
num_instances = len(masks) | |
if keypoints is not None: | |
if num_instances: | |
assert len(keypoints) == num_instances | |
else: | |
num_instances = len(keypoints) | |
keypoints = self._convert_keypoints(keypoints) | |
if labels is not None: | |
assert len(labels) == num_instances | |
if assigned_colors is None: | |
assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)] | |
if num_instances == 0: | |
return self.output | |
if boxes is not None and boxes.shape[1] == 5: | |
return self.overlay_rotated_instances( | |
boxes=boxes, labels=labels, assigned_colors=assigned_colors | |
) | |
# Display in largest to smallest order to reduce occlusion. | |
areas = None | |
if boxes is not None: | |
areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1) | |
elif masks is not None: | |
areas = np.asarray([x.area() for x in masks]) | |
if areas is not None: | |
sorted_idxs = np.argsort(-areas).tolist() | |
# Re-order overlapped instances in descending order. | |
boxes = boxes[sorted_idxs] if boxes is not None else None | |
labels = [labels[k] for k in sorted_idxs] if labels is not None else None | |
masks = [masks[idx] for idx in sorted_idxs] if masks is not None else None | |
rec = [rec[idx] for idx in sorted_idxs] if rec is not None else None | |
# rec_score = [rec_score[idx] for idx in sorted_idxs] if rec is not None else None | |
scores = [scores[idx] for idx in sorted_idxs] if scores is not None else None | |
# assigned_colors = [assigned_colors[idx] for idx in sorted_idxs] | |
keypoints = keypoints[sorted_idxs] if keypoints is not None else None | |
#luyao# | |
assigned_colors = [[0,113.985,118.955],[216.75,82.875,24.99],[236.895, 176.97, 31.875],[125.97, 46.92, 141.78],[118.83, 171.87, 47.94],[76.755, 189.975, 237.915],[161.925, 19.89, 46.92],\ | |
[255,140,0 ],[70,130,180 ],[128,128,0 ],[205,92,92 ],[128,0,128 ],[255,182,193],[255,255,0],[105,105,105],[0,255,255],[0,255,0 ],\ | |
[210,180,140],[255,0,0 ],[0,139,139],[255,0,255],[127,255,0],[75,0,130],[32,178,170],[255,215,0],[219,112,147],[148,0,211 ],\ | |
[100,149,237],[175,238,238 ],[143,188,143],[255,255,224 ],[244,164,96],[188,143,143],[192,192,192 ],[220,20,60],[218,112,214],[147,112,219]] | |
rec = [_ctc_decode_recognition(rrec) for rrec in rec] | |
# assigned_colors = [[1,140/255,0],[30/255,144/255,1],[148/255,0,211/255],[0,1,1],[1,0,0],\ | |
# [30/255,143/255,1],[0.94,0.5,0.5],[1,1,0],[0.5,0.5,0],[0.823,0.412,0.117],[0.58,0,0.827],[0.5,0,0]\ | |
# ,[0.82,0.41,0.12],[0.41,0.41,0.41],[0,0.54,0.54],[0.75,0.25,0.65],[0.2,0.6,0.8],[0.74,0,0.3],[0,1.0,0.4],[1,0.5,0.5],[0.5,0.5,1]\ | |
# ,[0.6,0,1],[0.56,0.56,0.3],[0,1,0],[1.0,0.0,0.4],[0.0,1.0,0.4],[0.0,0.5,1.0],[1,215/255,0]] | |
poly = [] | |
alpha = 0.4 | |
for i in range(num_instances): | |
if masks is not None: | |
poly.append(masks[i].polygons[0].astype(int).reshape(-1,2)) | |
keep = py_cpu_pnms(poly,scores,0.5) | |
for i in range(num_instances): | |
# if rec[i] == ' ': | |
# continue | |
if i not in keep: | |
continue | |
# color = assigned_colors[i] | |
# print(i) | |
color_ = assigned_colors[i%len(assigned_colors)] | |
color = [x/255 for x in color_] | |
# if boxes is not None: | |
# self.draw_box(boxes[i], edge_color=color) | |
#luyao | |
# alpha = 0.6 | |
H, W, _ = self.img.shape | |
if masks is not None: | |
for segment in masks[i].polygons: | |
segment = polygon2rbox(segment, H, W) | |
segment = np.array(segment) | |
self.draw_polygon(segment.reshape(-1, 2), color, alpha=alpha) | |
if labels is not None: | |
# first get a box | |
if boxes is not None: | |
#luyao# | |
x0, y0, x1, y1 = boxes[i] | |
text_pos = (x0, y0) # if drawing boxes, put text on the box corner. | |
horiz_align = "left" | |
elif masks is not None: | |
# skip small mask without polygon | |
if len(masks[i].polygons) == 0: | |
continue | |
x0, y0, x1, y1 = masks[i].bbox() | |
# draw text in the center (defined by median) when box is not drawn | |
# median is less sensitive to outliers. | |
text_pos = np.median(masks[i].mask.nonzero(), axis=1)[::-1] | |
horiz_align = "center" | |
else: | |
continue # drawing the box confidence for keypoints isn't very useful. | |
# for small objects, draw text at the side to avoid occlusion | |
instance_area = (y1 - y0) * (x1 - x0) | |
# print(x0,' ',x1,' ',y0,' ',y1,' ',self.output.height,' ', self.output.width) | |
#luyao# | |
if y0<5: | |
text_pos = ((x0+x1)//2,(y0+y1)//2) | |
#luyao# | |
# if ( | |
# instance_area < _SMALL_OBJECT_AREA_THRESH * self.output.scale | |
# or y1 - y0 < 40 * self.output.scale | |
# ): | |
# if y1 >= self.output.height - 5: | |
# text_pos = (x1, y0) | |
# else: | |
# text_pos = (x0, y1) | |
height_ratio = (y1 - y0) / np.sqrt(self.output.height * self.output.width) | |
lighter_color = self._change_color_brightness(color, brightness_factor=0.7) | |
font_size = ( | |
np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) | |
* 1.0 | |
* self._default_font_size | |
) | |
self.draw_text( | |
# labels[i], | |
# '', | |
rec[i], | |
text_pos, | |
color=lighter_color, | |
horizontal_alignment=horiz_align, | |
font_size=font_size, | |
) | |
# draw keypoints | |
if keypoints is not None: | |
for keypoints_per_instance in keypoints: | |
self.draw_and_connect_keypoints(keypoints_per_instance) | |
return self.output | |
def overlay_rotated_instances(self, boxes=None, labels=None, assigned_colors=None): | |
""" | |
Args: | |
boxes (ndarray): an Nx5 numpy array of | |
(x_center, y_center, width, height, angle_degrees) format | |
for the N objects in a single image. | |
labels (list[str]): the text to be displayed for each instance. | |
assigned_colors (list[matplotlib.colors]): a list of colors, where each color | |
corresponds to each mask or box in the image. Refer to 'matplotlib.colors' | |
for full list of formats that the colors are accepted in. | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
num_instances = len(boxes) | |
if assigned_colors is None: | |
assigned_colors = [random_color(rgb=True, maximum=1) for _ in range(num_instances)] | |
if num_instances == 0: | |
return self.output | |
# Display in largest to smallest order to reduce occlusion. | |
if boxes is not None: | |
areas = boxes[:, 2] * boxes[:, 3] | |
sorted_idxs = np.argsort(-areas).tolist() | |
# Re-order overlapped instances in descending order. | |
boxes = boxes[sorted_idxs] | |
labels = [labels[k] for k in sorted_idxs] if labels is not None else None | |
colors = [assigned_colors[idx] for idx in sorted_idxs] | |
for i in range(num_instances): | |
self.draw_rotated_box_with_label( | |
boxes[i], edge_color=colors[i], label=labels[i] if labels is not None else None | |
) | |
return self.output | |
def draw_and_connect_keypoints(self, keypoints): | |
""" | |
Draws keypoints of an instance and follows the rules for keypoint connections | |
to draw lines between appropriate keypoints. This follows color heuristics for | |
line color. | |
Args: | |
keypoints (Tensor): a tensor of shape (K, 3), where K is the number of keypoints | |
and the last dimension corresponds to (x, y, probability). | |
Returns: | |
output (VisImage): image object with visualizations. | |
""" | |
visible = {} | |
keypoint_names = self.metadata.get("keypoint_names") | |
for idx, keypoint in enumerate(keypoints): | |
# draw keypoint | |
x, y, prob = keypoint | |
if prob > _KEYPOINT_THRESHOLD: | |
self.draw_circle((x, y), color=_RED) | |
if keypoint_names: | |
keypoint_name = keypoint_names[idx] | |
visible[keypoint_name] = (x, y) | |
if self.metadata.get("keypoint_connection_rules"): | |
for kp0, kp1, color in self.metadata.keypoint_connection_rules: | |
if kp0 in visible and kp1 in visible: | |
x0, y0 = visible[kp0] | |
x1, y1 = visible[kp1] | |
color = tuple(x / 255.0 for x in color) | |
self.draw_line([x0, x1], [y0, y1], color=color) | |
# draw lines from nose to mid-shoulder and mid-shoulder to mid-hip | |
# Note that this strategy is specific to person keypoints. | |
# For other keypoints, it should just do nothing | |
try: | |
ls_x, ls_y = visible["left_shoulder"] | |
rs_x, rs_y = visible["right_shoulder"] | |
mid_shoulder_x, mid_shoulder_y = (ls_x + rs_x) / 2, (ls_y + rs_y) / 2 | |
except KeyError: | |
pass | |
else: | |
# draw line from nose to mid-shoulder | |
nose_x, nose_y = visible.get("nose", (None, None)) | |
if nose_x is not None: | |
self.draw_line([nose_x, mid_shoulder_x], [nose_y, mid_shoulder_y], color=_RED) | |
try: | |
# draw line from mid-shoulder to mid-hip | |
lh_x, lh_y = visible["left_hip"] | |
rh_x, rh_y = visible["right_hip"] | |
except KeyError: | |
pass | |
else: | |
mid_hip_x, mid_hip_y = (lh_x + rh_x) / 2, (lh_y + rh_y) / 2 | |
self.draw_line([mid_hip_x, mid_shoulder_x], [mid_hip_y, mid_shoulder_y], color=_RED) | |
return self.output | |
""" | |
Primitive drawing functions: | |
""" | |
def draw_text( | |
self, | |
text, | |
position, | |
*, | |
font_size=None, | |
color="g", | |
horizontal_alignment="center", | |
rotation=0 | |
): | |
""" | |
Args: | |
text (str): class label | |
position (tuple): a tuple of the x and y coordinates to place text on image. | |
font_size (int, optional): font of the text. If not provided, a font size | |
proportional to the image width is calculated and used. | |
color: color of the text. Refer to `matplotlib.colors` for full list | |
of formats that are accepted. | |
horizontal_alignment (str): see `matplotlib.text.Text` | |
rotation: rotation angle in degrees CCW | |
Returns: | |
output (VisImage): image object with text drawn. | |
""" | |
if not font_size: | |
font_size = self._default_font_size | |
# print(font_size, self.output.scale) | |
# since the text background is dark, we don't want the text to be dark | |
# color = np.maximum(list(mplc.to_rgb(color)), 0.2) | |
# color[np.argmax(color)] = max(0.8, np.max(color)) | |
#luyao# | |
color = 'w' | |
# font_size = 7.0 | |
x, y = position | |
font_path = "simsun.ttc" | |
prop = mfm.FontProperties(fname=font_path) | |
self.output.ax.text( | |
x, | |
y, | |
text, | |
size=font_size * self.output.scale, | |
# family="sans-serif", | |
family="monospace", | |
# family="serif", | |
#luyao# | |
bbox={"facecolor": "black", "alpha": 0.0, "pad": 0.0, "edgecolor": "none"}, | |
# bbox={"facecolor": "black", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"}, | |
# verticalalignment="top", | |
verticalalignment="bottom", | |
horizontalalignment=horizontal_alignment, | |
color=color, | |
zorder=10, | |
rotation=rotation, | |
fontproperties=prop, | |
#luyao | |
# fontweight='light' | |
) | |
return self.output | |
def draw_box(self, box_coord, alpha=0.5, edge_color="g", line_style="-"): | |
""" | |
Args: | |
box_coord (tuple): a tuple containing x0, y0, x1, y1 coordinates, where x0 and y0 | |
are the coordinates of the image's top left corner. x1 and y1 are the | |
coordinates of the image's bottom right corner. | |
alpha (float): blending efficient. Smaller values lead to more transparent masks. | |
edge_color: color of the outline of the box. Refer to `matplotlib.colors` | |
for full list of formats that are accepted. | |
line_style (string): the string to use to create the outline of the boxes. | |
Returns: | |
output (VisImage): image object with box drawn. | |
""" | |
x0, y0, x1, y1 = box_coord | |
width = x1 - x0 | |
height = y1 - y0 | |
# linewidth = max(self._default_font_size / 16, 1) | |
# linewidth = max(self._default_font_size / 4, 1) | |
#luyao# | |
edge_color=[0.196,0.80,0.196] | |
alpha = 1.0 | |
linewidth = 0.7 | |
self.output.ax.add_patch( | |
mpl.patches.Rectangle( | |
(x0, y0), | |
width, | |
height, | |
fill=False, | |
edgecolor=edge_color, | |
linewidth=linewidth * self.output.scale, | |
alpha=alpha, | |
linestyle=line_style, | |
) | |
) | |
return self.output | |
def draw_rotated_box_with_label( | |
self, rotated_box, alpha=0.5, edge_color="g", line_style="-", label=None | |
): | |
""" | |
Draw a rotated box with label on its top-left corner. | |
Args: | |
rotated_box (tuple): a tuple containing (cnt_x, cnt_y, w, h, angle), | |
where cnt_x and cnt_y are the center coordinates of the box. | |
w and h are the width and height of the box. angle represents how | |
many degrees the box is rotated CCW with regard to the 0-degree box. | |
alpha (float): blending efficient. Smaller values lead to more transparent masks. | |
edge_color: color of the outline of the box. Refer to `matplotlib.colors` | |
for full list of formats that are accepted. | |
line_style (string): the string to use to create the outline of the boxes. | |
label (string): label for rotated box. It will not be rendered when set to None. | |
Returns: | |
output (VisImage): image object with box drawn. | |
""" | |
cnt_x, cnt_y, w, h, angle = rotated_box | |
area = w * h | |
# use thinner lines when the box is small | |
linewidth = self._default_font_size / ( | |
6 if area < _SMALL_OBJECT_AREA_THRESH * self.output.scale else 3 | |
) | |
theta = angle * math.pi / 180.0 | |
c = math.cos(theta) | |
s = math.sin(theta) | |
rect = [(-w / 2, h / 2), (-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2)] | |
# x: left->right ; y: top->down | |
rotated_rect = [(s * yy + c * xx + cnt_x, c * yy - s * xx + cnt_y) for (xx, yy) in rect] | |
for k in range(4): | |
j = (k + 1) % 4 | |
self.draw_line( | |
[rotated_rect[k][0], rotated_rect[j][0]], | |
[rotated_rect[k][1], rotated_rect[j][1]], | |
color=edge_color, | |
linestyle="--" if k == 1 else line_style, | |
linewidth=linewidth, | |
) | |
if label is not None: | |
text_pos = rotated_rect[1] # topleft corner | |
height_ratio = h / np.sqrt(self.output.height * self.output.width) | |
label_color = self._change_color_brightness(edge_color, brightness_factor=0.7) | |
font_size = ( | |
np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) * 0.5 * self._default_font_size | |
) | |
self.draw_text(label, text_pos, color=label_color, font_size=font_size, rotation=angle) | |
return self.output | |
def draw_circle(self, circle_coord, color, radius=3): | |
""" | |
Args: | |
circle_coord (list(int) or tuple(int)): contains the x and y coordinates | |
of the center of the circle. | |
color: color of the polygon. Refer to `matplotlib.colors` for a full list of | |
formats that are accepted. | |
radius (int): radius of the circle. | |
Returns: | |
output (VisImage): image object with box drawn. | |
""" | |
x, y = circle_coord | |
self.output.ax.add_patch( | |
mpl.patches.Circle(circle_coord, radius=radius, fill=True, color=color) | |
) | |
return self.output | |
def draw_line(self, x_data, y_data, color, linestyle="-", linewidth=None): | |
""" | |
Args: | |
x_data (list[int]): a list containing x values of all the points being drawn. | |
Length of list should match the length of y_data. | |
y_data (list[int]): a list containing y values of all the points being drawn. | |
Length of list should match the length of x_data. | |
color: color of the line. Refer to `matplotlib.colors` for a full list of | |
formats that are accepted. | |
linestyle: style of the line. Refer to `matplotlib.lines.Line2D` | |
for a full list of formats that are accepted. | |
linewidth (float or None): width of the line. When it's None, | |
a default value will be computed and used. | |
Returns: | |
output (VisImage): image object with line drawn. | |
""" | |
if linewidth is None: | |
linewidth = self._default_font_size / 3 | |
linewidth = max(linewidth, 1) | |
self.output.ax.add_line( | |
mpl.lines.Line2D( | |
x_data, | |
y_data, | |
linewidth=linewidth * self.output.scale, | |
color=color, | |
linestyle=linestyle, | |
) | |
) | |
return self.output | |
def draw_binary_mask( | |
self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=0 | |
): | |
""" | |
Args: | |
binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and | |
W is the image width. Each value in the array is either a 0 or 1 value of uint8 | |
type. | |
color: color of the mask. Refer to `matplotlib.colors` for a full list of | |
formats that are accepted. If None, will pick a random color. | |
edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a | |
full list of formats that are accepted. | |
text (str): if None, will be drawn in the object's center of mass. | |
alpha (float): blending efficient. Smaller values lead to more transparent masks. | |
area_threshold (float): a connected component small than this will not be shown. | |
Returns: | |
output (VisImage): image object with mask drawn. | |
""" | |
if color is None: | |
color = random_color(rgb=True, maximum=1) | |
color = mplc.to_rgb(color) | |
has_valid_segment = False | |
binary_mask = binary_mask.astype("uint8") # opencv needs uint8 | |
mask = GenericMask(binary_mask, self.output.height, self.output.width) | |
shape2d = (binary_mask.shape[0], binary_mask.shape[1]) | |
if not mask.has_holes: | |
# draw polygons for regular masks | |
for segment in mask.polygons: | |
area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1])) | |
if area < (area_threshold or 0): | |
continue | |
has_valid_segment = True | |
segment = segment.reshape(-1, 2) | |
self.draw_polygon(segment, color=color, edge_color=edge_color, alpha=alpha) | |
else: | |
# TODO: Use Path/PathPatch to draw vector graphics: | |
# https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon | |
rgba = np.zeros(shape2d + (4,), dtype="float32") | |
rgba[:, :, :3] = color | |
rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha | |
has_valid_segment = True | |
self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0)) | |
if text is not None and has_valid_segment: | |
# TODO sometimes drawn on wrong objects. the heuristics here can improve. | |
lighter_color = self._change_color_brightness(color, brightness_factor=0.7) | |
_num_cc, cc_labels, stats, centroids = cv2.connectedComponentsWithStats(binary_mask, 8) | |
largest_component_id = np.argmax(stats[1:, -1]) + 1 | |
# draw text on the largest component, as well as other very large components. | |
for cid in range(1, _num_cc): | |
if cid == largest_component_id or stats[cid, -1] > _LARGE_MASK_AREA_THRESH: | |
# median is more stable than centroid | |
# center = centroids[largest_component_id] | |
center = np.median((cc_labels == cid).nonzero(), axis=1)[::-1] | |
self.draw_text(text, center, color=lighter_color) | |
return self.output | |
def draw_polygon(self, segment, color, edge_color=None, alpha=0.5): | |
""" | |
Args: | |
segment: numpy array of shape Nx2, containing all the points in the polygon. | |
color: color of the polygon. Refer to `matplotlib.colors` for a full list of | |
formats that are accepted. | |
edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a | |
full list of formats that are accepted. If not provided, a darker shade | |
of the polygon color will be used instead. | |
alpha (float): blending efficient. Smaller values lead to more transparent masks. | |
Returns: | |
output (VisImage): image object with polygon drawn. | |
""" | |
#luyao# | |
# edge_color = [] | |
if edge_color is None: | |
# make edge color darker than the polygon color | |
if alpha > 0.8: | |
edge_color = self._change_color_brightness(color, brightness_factor=-0.7) | |
else: | |
edge_color = color | |
edge_color = mplc.to_rgb(edge_color) + (1,) | |
polygon = mpl.patches.Polygon( | |
segment, | |
fill=True, | |
facecolor=mplc.to_rgb(color) + (alpha,), | |
#luyao# qudiaomaskyanse | |
# edgecolor=edge_color, | |
linewidth=max(self._default_font_size // 15 * self.output.scale, 1), | |
) | |
self.output.ax.add_patch(polygon) | |
return self.output | |
""" | |
Internal methods: | |
""" | |
def _jitter(self, color): | |
""" | |
Randomly modifies given color to produce a slightly different color than the color given. | |
Args: | |
color (tuple[double]): a tuple of 3 elements, containing the RGB values of the color | |
picked. The values in the list are in the [0.0, 1.0] range. | |
Returns: | |
jittered_color (tuple[double]): a tuple of 3 elements, containing the RGB values of the | |
color after being jittered. The values in the list are in the [0.0, 1.0] range. | |
""" | |
color = mplc.to_rgb(color) | |
vec = np.random.rand(3) | |
# better to do it in another color space | |
vec = vec / np.linalg.norm(vec) * 0.5 | |
res = np.clip(vec + color, 0, 1) | |
return tuple(res) | |
def _create_grayscale_image(self, mask=None): | |
""" | |
Create a grayscale version of the original image. | |
The colors in masked area, if given, will be kept. | |
""" | |
img_bw = self.img.astype("f4").mean(axis=2) | |
img_bw = np.stack([img_bw] * 3, axis=2) | |
if mask is not None: | |
img_bw[mask] = self.img[mask] | |
return img_bw | |
def _change_color_brightness(self, color, brightness_factor): | |
""" | |
Depending on the brightness_factor, gives a lighter or darker color i.e. a color with | |
less or more saturation than the original color. | |
Args: | |
color: color of the polygon. Refer to `matplotlib.colors` for a full list of | |
formats that are accepted. | |
brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of | |
0 will correspond to no change, a factor in [-1.0, 0) range will result in | |
a darker color and a factor in (0, 1.0] range will result in a lighter color. | |
Returns: | |
modified_color (tuple[double]): a tuple containing the RGB values of the | |
modified color. Each value in the tuple is in the [0.0, 1.0] range. | |
""" | |
assert brightness_factor >= -1.0 and brightness_factor <= 1.0 | |
color = mplc.to_rgb(color) | |
polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color)) | |
modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1]) | |
modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness | |
modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness | |
modified_color = colorsys.hls_to_rgb(polygon_color[0], modified_lightness, polygon_color[2]) | |
return modified_color | |
def _convert_boxes(self, boxes): | |
""" | |
Convert different format of boxes to an NxB array, where B = 4 or 5 is the box dimension. | |
""" | |
if isinstance(boxes, Boxes) or isinstance(boxes, RotatedBoxes): | |
return boxes.tensor.numpy() | |
else: | |
return np.asarray(boxes) | |
def _convert_masks(self, masks_or_polygons): | |
""" | |
Convert different format of masks or polygons to a tuple of masks and polygons. | |
Returns: | |
list[GenericMask]: | |
""" | |
m = masks_or_polygons | |
if isinstance(m, PolygonMasks): | |
m = m.polygons | |
if isinstance(m, BitMasks): | |
m = m.tensor.numpy() | |
if isinstance(m, torch.Tensor): | |
m = m.numpy() | |
ret = [] | |
for x in m: | |
if isinstance(x, GenericMask): | |
ret.append(x) | |
else: | |
ret.append(GenericMask(x, self.output.height, self.output.width)) | |
return ret | |
def _convert_keypoints(self, keypoints): | |
if isinstance(keypoints, Keypoints): | |
keypoints = keypoints.tensor | |
keypoints = np.asarray(keypoints) | |
return keypoints | |
def get_output(self): | |
""" | |
Returns: | |
output (VisImage): the image output containing the visualizations added | |
to the image. | |
""" | |
return self.output | |
def polygon2rbox(polygon, image_height, image_width): | |
poly = np.array(polygon).reshape((-1, 2)).astype(np.float32) | |
rect = cv2.minAreaRect(poly) | |
corners = cv2.boxPoints(rect) | |
corners = np.array(corners, dtype="int") | |
pts = get_tight_rect(corners, 0, 0, image_height, image_width, 1) | |
pts = list(map(int, pts)) | |
return pts | |
def get_tight_rect(points, start_x, start_y, image_height, image_width, scale): | |
points = list(points) | |
ps = sorted(points, key=lambda x: x[0]) | |
if ps[1][1] > ps[0][1]: | |
px1 = ps[0][0] * scale + start_x | |
py1 = ps[0][1] * scale + start_y | |
px4 = ps[1][0] * scale + start_x | |
py4 = ps[1][1] * scale + start_y | |
else: | |
px1 = ps[1][0] * scale + start_x | |
py1 = ps[1][1] * scale + start_y | |
px4 = ps[0][0] * scale + start_x | |
py4 = ps[0][1] * scale + start_y | |
if ps[3][1] > ps[2][1]: | |
px2 = ps[2][0] * scale + start_x | |
py2 = ps[2][1] * scale + start_y | |
px3 = ps[3][0] * scale + start_x | |
py3 = ps[3][1] * scale + start_y | |
else: | |
px2 = ps[3][0] * scale + start_x | |
py2 = ps[3][1] * scale + start_y | |
px3 = ps[2][0] * scale + start_x | |
py3 = ps[2][1] * scale + start_y | |
px1 = min(max(px1, 1), image_width - 1) | |
px2 = min(max(px2, 1), image_width - 1) | |
px3 = min(max(px3, 1), image_width - 1) | |
px4 = min(max(px4, 1), image_width - 1) | |
py1 = min(max(py1, 1), image_height - 1) | |
py2 = min(max(py2, 1), image_height - 1) | |
py3 = min(max(py3, 1), image_height - 1) | |
py4 = min(max(py4, 1), image_height - 1) | |
return [px1, py1, px2, py2, px3, py3, px4, py4] | |