| 493 | |
| 494 | |
| 495 | class VisImage: |
| 496 | def __init__(self, img, scale=1.0): |
| 497 | self.img = img |
| 498 | self.scale = scale |
| 499 | self.width, self.height = img.shape[1], img.shape[0] |
| 500 | self._setup_figure(img) |
| 501 | |
| 502 | def _setup_figure(self, img): |
| 503 | fig = mplfigure.Figure(frameon=False) |
| 504 | self.dpi = fig.get_dpi() |
| 505 | # add a small 1e-2 to avoid precision lost due to matplotlib's truncation |
| 506 | # (https://github.com/matplotlib/matplotlib/issues/15363) |
| 507 | fig.set_size_inches( |
| 508 | (self.width * self.scale + 1e-2) / self.dpi, |
| 509 | (self.height * self.scale + 1e-2) / self.dpi, |
| 510 | ) |
| 511 | self.canvas = FigureCanvasAgg(fig) |
| 512 | # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig) |
| 513 | ax = fig.add_axes([0.0, 0.0, 1.0, 1.0]) |
| 514 | ax.axis("off") |
| 515 | self.fig = fig |
| 516 | self.ax = ax |
| 517 | self.reset_image(img) |
| 518 | |
| 519 | def reset_image(self, img): |
| 520 | img = img.astype("uint8") |
| 521 | self.ax.imshow(img, extent=(0, self.width, self.height, 0), interpolation="nearest") |
| 522 | |
| 523 | def save(self, filepath): |
| 524 | self.fig.savefig(filepath) |
| 525 | |
| 526 | def get_image(self): |
| 527 | canvas = self.canvas |
| 528 | s, (width, height) = canvas.print_to_buffer() |
| 529 | |
| 530 | buffer = np.frombuffer(s, dtype="uint8") |
| 531 | |
| 532 | img_rgba = buffer.reshape(height, width, 4) |
| 533 | rgb, alpha = np.split(img_rgba, [3], axis=2) |
| 534 | return rgb.astype("uint8") |
| 535 | |
| 536 | |
| 537 | class Visualizer: |