| 242 | |
| 243 | |
| 244 | class VisImage: |
| 245 | def __init__(self, img, scale=1.0): |
| 246 | """ |
| 247 | Args: |
| 248 | img (ndarray): an RGB image of shape (H, W, 3). |
| 249 | scale (float): scale the input image |
| 250 | """ |
| 251 | self.img = img |
| 252 | self.scale = scale |
| 253 | self.width, self.height = img.shape[1], img.shape[0] |
| 254 | self._setup_figure(img) |
| 255 | |
| 256 | def _setup_figure(self, img): |
| 257 | """ |
| 258 | Args: |
| 259 | Same as in :meth:`__init__()`. |
| 260 | |
| 261 | Returns: |
| 262 | fig (matplotlib.pyplot.figure): top level container for all the image plot elements. |
| 263 | ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system. |
| 264 | """ |
| 265 | fig = mplfigure.Figure(frameon=False) |
| 266 | self.dpi = fig.get_dpi() |
| 267 | # add a small 1e-2 to avoid precision lost due to matplotlib's truncation |
| 268 | # (https://github.com/matplotlib/matplotlib/issues/15363) |
| 269 | fig.set_size_inches( |
| 270 | (self.width * self.scale + 1e-2) / self.dpi, |
| 271 | (self.height * self.scale + 1e-2) / self.dpi, |
| 272 | ) |
| 273 | self.canvas = FigureCanvasAgg(fig) |
| 274 | # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig) |
| 275 | ax = fig.add_axes([0.0, 0.0, 1.0, 1.0]) |
| 276 | ax.axis("off") |
| 277 | # Need to imshow this first so that other patches can be drawn on top |
| 278 | ax.imshow(img, extent=(0, self.width, self.height, 0), interpolation="nearest") |
| 279 | |
| 280 | self.fig = fig |
| 281 | self.ax = ax |
| 282 | |
| 283 | def save(self, filepath): |
| 284 | """ |
| 285 | Args: |
| 286 | filepath (str): a string that contains the absolute path, including the file name, where |
| 287 | the visualized image will be saved. |
| 288 | """ |
| 289 | self.fig.savefig(filepath) |
| 290 | |
| 291 | def get_image(self): |
| 292 | """ |
| 293 | Returns: |
| 294 | ndarray: |
| 295 | the visualized image of shape (H, W, 3) (RGB) in uint8 type. |
| 296 | The shape is scaled w.r.t the input image using the given `scale` argument. |
| 297 | """ |
| 298 | canvas = self.canvas |
| 299 | s, (width, height) = canvas.print_to_buffer() |
| 300 | # buf = io.BytesIO() # works for cairo backend |
| 301 | # canvas.print_rgba(buf) |