| """ |
| coding=utf-8 |
| Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal |
| Adapted From Facebook Inc, Detectron2 |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| http://www.apache.org/licenses/LICENSE-2.0 |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License.import copy |
| """ |
| import colorsys |
| import io |
|
|
| import matplotlib as mpl |
| import matplotlib.colors as mplc |
| import matplotlib.figure as mplfigure |
| import numpy as np |
| import torch |
| from matplotlib.backends.backend_agg import FigureCanvasAgg |
|
|
| import cv2 |
| from .utils import img_tensorize |
|
|
|
|
| _SMALL_OBJ = 1000 |
|
|
|
|
| class SingleImageViz: |
| def __init__( |
| self, |
| img, |
| scale=1.2, |
| edgecolor="g", |
| alpha=0.5, |
| linestyle="-", |
| saveas="test_out.jpg", |
| rgb=True, |
| pynb=False, |
| id2obj=None, |
| id2attr=None, |
| pad=0.7, |
| ): |
| """ |
| img: an RGB image of shape (H, W, 3). |
| """ |
| if isinstance(img, torch.Tensor): |
| img = img.numpy().astype("np.uint8") |
| if isinstance(img, str): |
| img = img_tensorize(img) |
| assert isinstance(img, np.ndarray) |
|
|
| width, height = img.shape[1], img.shape[0] |
| fig = mplfigure.Figure(frameon=False) |
| dpi = fig.get_dpi() |
| width_in = (width * scale + 1e-2) / dpi |
| height_in = (height * scale + 1e-2) / dpi |
| fig.set_size_inches(width_in, height_in) |
| ax = fig.add_axes([0.0, 0.0, 1.0, 1.0]) |
| ax.axis("off") |
| ax.set_xlim(0.0, width) |
| ax.set_ylim(height) |
|
|
| self.saveas = saveas |
| self.rgb = rgb |
| self.pynb = pynb |
| self.img = img |
| self.edgecolor = edgecolor |
| self.alpha = 0.5 |
| self.linestyle = linestyle |
| self.font_size = int(np.sqrt(min(height, width)) * scale // 3) |
| self.width = width |
| self.height = height |
| self.scale = scale |
| self.fig = fig |
| self.ax = ax |
| self.pad = pad |
| self.id2obj = id2obj |
| self.id2attr = id2attr |
| self.canvas = FigureCanvasAgg(fig) |
|
|
| def add_box(self, box, color=None): |
| if color is None: |
| color = self.edgecolor |
| (x0, y0, x1, y1) = box |
| width = x1 - x0 |
| height = y1 - y0 |
| self.ax.add_patch( |
| mpl.patches.Rectangle( |
| (x0, y0), |
| width, |
| height, |
| fill=False, |
| edgecolor=color, |
| linewidth=self.font_size // 3, |
| alpha=self.alpha, |
| linestyle=self.linestyle, |
| ) |
| ) |
|
|
| def draw_boxes(self, boxes, obj_ids=None, obj_scores=None, attr_ids=None, attr_scores=None): |
| if len(boxes.shape) > 2: |
| boxes = boxes[0] |
| if len(obj_ids.shape) > 1: |
| obj_ids = obj_ids[0] |
| if len(obj_scores.shape) > 1: |
| obj_scores = obj_scores[0] |
| if len(attr_ids.shape) > 1: |
| attr_ids = attr_ids[0] |
| if len(attr_scores.shape) > 1: |
| attr_scores = attr_scores[0] |
| if isinstance(boxes, torch.Tensor): |
| boxes = boxes.numpy() |
| if isinstance(boxes, list): |
| boxes = np.array(boxes) |
| assert isinstance(boxes, np.ndarray) |
| areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1) |
| sorted_idxs = np.argsort(-areas).tolist() |
| boxes = boxes[sorted_idxs] if boxes is not None else None |
| obj_ids = obj_ids[sorted_idxs] if obj_ids is not None else None |
| obj_scores = obj_scores[sorted_idxs] if obj_scores is not None else None |
| attr_ids = attr_ids[sorted_idxs] if attr_ids is not None else None |
| attr_scores = attr_scores[sorted_idxs] if attr_scores is not None else None |
|
|
| assigned_colors = [self._random_color(maximum=1) for _ in range(len(boxes))] |
| assigned_colors = [assigned_colors[idx] for idx in sorted_idxs] |
| if obj_ids is not None: |
| labels = self._create_text_labels_attr(obj_ids, obj_scores, attr_ids, attr_scores) |
| for i in range(len(boxes)): |
| color = assigned_colors[i] |
| self.add_box(boxes[i], color) |
| self.draw_labels(labels[i], boxes[i], color) |
|
|
| def draw_labels(self, label, box, color): |
| x0, y0, x1, y1 = box |
| text_pos = (x0, y0) |
| instance_area = (y1 - y0) * (x1 - x0) |
| small = _SMALL_OBJ * self.scale |
| if instance_area < small or y1 - y0 < 40 * self.scale: |
| if y1 >= self.height - 5: |
| text_pos = (x1, y0) |
| else: |
| text_pos = (x0, y1) |
|
|
| height_ratio = (y1 - y0) / np.sqrt(self.height * self.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) |
| font_size *= 0.75 * self.font_size |
|
|
| self.draw_text( |
| text=label, |
| position=text_pos, |
| color=lighter_color, |
| ) |
|
|
| def draw_text( |
| self, |
| text, |
| position, |
| color="g", |
| ha="left", |
| ): |
| rotation = 0 |
| font_size = self.font_size |
| color = np.maximum(list(mplc.to_rgb(color)), 0.2) |
| color[np.argmax(color)] = max(0.8, np.max(color)) |
| bbox = { |
| "facecolor": "black", |
| "alpha": self.alpha, |
| "pad": self.pad, |
| "edgecolor": "none", |
| } |
| x, y = position |
| self.ax.text( |
| x, |
| y, |
| text, |
| size=font_size * self.scale, |
| family="sans-serif", |
| bbox=bbox, |
| verticalalignment="top", |
| horizontalalignment=ha, |
| color=color, |
| zorder=10, |
| rotation=rotation, |
| ) |
|
|
| def save(self, saveas=None): |
| if saveas is None: |
| saveas = self.saveas |
| if saveas.lower().endswith(".jpg") or saveas.lower().endswith(".png"): |
| cv2.imwrite( |
| saveas, |
| self._get_buffer()[:, :, ::-1], |
| ) |
| else: |
| self.fig.savefig(saveas) |
|
|
| def _create_text_labels_attr(self, classes, scores, attr_classes, attr_scores): |
| labels = [self.id2obj[i] for i in classes] |
| attr_labels = [self.id2attr[i] for i in attr_classes] |
| labels = [ |
| f"{label} {score:.2f} {attr} {attr_score:.2f}" |
| for label, score, attr, attr_score in zip(labels, scores, attr_labels, attr_scores) |
| ] |
| return labels |
|
|
| def _create_text_labels(self, classes, scores): |
| labels = [self.id2obj[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(li, s * 100) for li, s in zip(labels, scores)] |
| return labels |
|
|
| def _random_color(self, maximum=255): |
| idx = np.random.randint(0, len(_COLORS)) |
| ret = _COLORS[idx] * maximum |
| if not self.rgb: |
| ret = ret[::-1] |
| return ret |
|
|
| def _get_buffer(self): |
| if not self.pynb: |
| s, (width, height) = self.canvas.print_to_buffer() |
| if (width, height) != (self.width, self.height): |
| img = cv2.resize(self.img, (width, height)) |
| else: |
| img = self.img |
| else: |
| buf = io.BytesIO() |
| self.canvas.print_rgba(buf) |
| width, height = self.width, self.height |
| s = buf.getvalue() |
| img = self.img |
|
|
| buffer = np.frombuffer(s, dtype="uint8") |
| img_rgba = buffer.reshape(height, width, 4) |
| rgb, alpha = np.split(img_rgba, [3], axis=2) |
|
|
| try: |
| import numexpr as ne |
|
|
| visualized_image = ne.evaluate("img * (1 - alpha / 255.0) + rgb * (alpha / 255.0)") |
| except ImportError: |
| alpha = alpha.astype("float32") / 255.0 |
| visualized_image = img * (1 - alpha) + rgb * alpha |
|
|
| return visualized_image.astype("uint8") |
|
|
| def _change_color_brightness(self, color, brightness_factor): |
| 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 |
|
|
|
|
| |
| _COLORS = ( |
| np.array( |
| [ |
| 0.000, |
| 0.447, |
| 0.741, |
| 0.850, |
| 0.325, |
| 0.098, |
| 0.929, |
| 0.694, |
| 0.125, |
| 0.494, |
| 0.184, |
| 0.556, |
| 0.466, |
| 0.674, |
| 0.188, |
| 0.301, |
| 0.745, |
| 0.933, |
| 0.635, |
| 0.078, |
| 0.184, |
| 0.300, |
| 0.300, |
| 0.300, |
| 0.600, |
| 0.600, |
| 0.600, |
| 1.000, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.500, |
| 0.000, |
| 0.749, |
| 0.749, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.000, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.667, |
| 0.000, |
| 1.000, |
| 0.333, |
| 0.333, |
| 0.000, |
| 0.333, |
| 0.667, |
| 0.000, |
| 0.333, |
| 1.000, |
| 0.000, |
| 0.667, |
| 0.333, |
| 0.000, |
| 0.667, |
| 0.667, |
| 0.000, |
| 0.667, |
| 1.000, |
| 0.000, |
| 1.000, |
| 0.333, |
| 0.000, |
| 1.000, |
| 0.667, |
| 0.000, |
| 1.000, |
| 1.000, |
| 0.000, |
| 0.000, |
| 0.333, |
| 0.500, |
| 0.000, |
| 0.667, |
| 0.500, |
| 0.000, |
| 1.000, |
| 0.500, |
| 0.333, |
| 0.000, |
| 0.500, |
| 0.333, |
| 0.333, |
| 0.500, |
| 0.333, |
| 0.667, |
| 0.500, |
| 0.333, |
| 1.000, |
| 0.500, |
| 0.667, |
| 0.000, |
| 0.500, |
| 0.667, |
| 0.333, |
| 0.500, |
| 0.667, |
| 0.667, |
| 0.500, |
| 0.667, |
| 1.000, |
| 0.500, |
| 1.000, |
| 0.000, |
| 0.500, |
| 1.000, |
| 0.333, |
| 0.500, |
| 1.000, |
| 0.667, |
| 0.500, |
| 1.000, |
| 1.000, |
| 0.500, |
| 0.000, |
| 0.333, |
| 1.000, |
| 0.000, |
| 0.667, |
| 1.000, |
| 0.000, |
| 1.000, |
| 1.000, |
| 0.333, |
| 0.000, |
| 1.000, |
| 0.333, |
| 0.333, |
| 1.000, |
| 0.333, |
| 0.667, |
| 1.000, |
| 0.333, |
| 1.000, |
| 1.000, |
| 0.667, |
| 0.000, |
| 1.000, |
| 0.667, |
| 0.333, |
| 1.000, |
| 0.667, |
| 0.667, |
| 1.000, |
| 0.667, |
| 1.000, |
| 1.000, |
| 1.000, |
| 0.000, |
| 1.000, |
| 1.000, |
| 0.333, |
| 1.000, |
| 1.000, |
| 0.667, |
| 1.000, |
| 0.333, |
| 0.000, |
| 0.000, |
| 0.500, |
| 0.000, |
| 0.000, |
| 0.667, |
| 0.000, |
| 0.000, |
| 0.833, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.000, |
| 0.000, |
| 0.000, |
| 0.167, |
| 0.000, |
| 0.000, |
| 0.333, |
| 0.000, |
| 0.000, |
| 0.500, |
| 0.000, |
| 0.000, |
| 0.667, |
| 0.000, |
| 0.000, |
| 0.833, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.000, |
| 0.000, |
| 0.000, |
| 0.167, |
| 0.000, |
| 0.000, |
| 0.333, |
| 0.000, |
| 0.000, |
| 0.500, |
| 0.000, |
| 0.000, |
| 0.667, |
| 0.000, |
| 0.000, |
| 0.833, |
| 0.000, |
| 0.000, |
| 1.000, |
| 0.000, |
| 0.000, |
| 0.000, |
| 0.143, |
| 0.143, |
| 0.143, |
| 0.857, |
| 0.857, |
| 0.857, |
| 1.000, |
| 1.000, |
| 1.000, |
| ] |
| ) |
| .astype(np.float32) |
| .reshape(-1, 3) |
| ) |
