Create visualizing_image.py

visualizing_image.py

@@ -0,0 +1,496 @@
+"""
+ 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()  # works for cairo backend
+            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  # fuse them with numexpr
+
+            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
+
+
+# Color map
+_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)
+)