Spaces:
Running
on
Zero
Running
on
Zero
| from contextlib import contextmanager | |
| import logging | |
| import os | |
| from matplotlib.pyplot import figimage, margins | |
| import numpy as np | |
| import cv2 | |
| try: | |
| import matplotlib.pyplot as plt # pylint: disable=import-error | |
| except ModuleNotFoundError as err: | |
| if err.name != 'matplotlib': | |
| raise err | |
| plt = None | |
| LOG = logging.getLogger(__name__) | |
| class Canvas: | |
| """Canvas for plotting. | |
| All methods expose Axes objects. To get Figure objects, you can ask the axis | |
| `ax.get_figure()`. | |
| """ | |
| all_images_directory = None | |
| all_images_count = 0 | |
| show = False | |
| image_width = 7.0 | |
| image_height = None | |
| blank_dpi = 200 | |
| image_dpi_factor = 1.0 | |
| image_min_dpi = 50.0 | |
| out_file_extension = 'pdf' | |
| white_overlay = False | |
| def generic_name(cls): | |
| if cls.all_images_directory is None: | |
| return None | |
| os.makedirs(cls.all_images_directory, exist_ok=True) | |
| cls.all_images_count += 1 | |
| return os.path.join(cls.all_images_directory, | |
| '{:04}.{}'.format(cls.all_images_count, cls.out_file_extension)) | |
| def blank(cls, fig_file=None, *, dpi=None, nomargin=False, **kwargs): | |
| if plt is None: | |
| raise Exception('please install matplotlib') | |
| if fig_file is None: | |
| fig_file = cls.generic_name() | |
| if dpi is None: | |
| dpi = cls.blank_dpi | |
| if 'figsize' not in kwargs: | |
| kwargs['figsize'] = (10, 6) | |
| if nomargin: | |
| if 'gridspec_kw' not in kwargs: | |
| kwargs['gridspec_kw'] = {} | |
| kwargs['gridspec_kw']['wspace'] = 0 | |
| kwargs['gridspec_kw']['hspace'] = 0 | |
| kwargs['gridspec_kw']['left'] = 0.0 | |
| kwargs['gridspec_kw']['right'] = 1.0 | |
| kwargs['gridspec_kw']['top'] = 1.0 | |
| kwargs['gridspec_kw']['bottom'] = 0.0 | |
| fig, ax = plt.subplots(dpi=dpi, **kwargs) | |
| yield ax | |
| fig.set_tight_layout(not margins) | |
| if fig_file: | |
| LOG.debug('writing image to %s', fig_file) | |
| fig.savefig(fig_file) | |
| if cls.show: | |
| plt.show() | |
| plt.close(fig) | |
| def image(cls, image, fig_file=None, *, margin=None, **kwargs): | |
| if plt is None: | |
| raise Exception('please install matplotlib') | |
| if fig_file is None: | |
| fig_file = cls.generic_name() | |
| if isinstance(image, str): | |
| image = cv2.imread(image)[...,::-1] | |
| else: | |
| image = np.asarray(image) | |
| if margin is None: | |
| margin = [0.0, 0.0, 0.0, 0.0] | |
| elif isinstance(margin, float): | |
| margin = [margin, margin, margin, margin] | |
| assert len(margin) == 4 | |
| if 'figsize' not in kwargs: | |
| # compute figure size: use image ratio and take the drawable area | |
| # into account that is left after subtracting margins. | |
| image_ratio = image.shape[0] / image.shape[1] | |
| image_area_ratio = (1.0 - margin[1] - margin[3]) / (1.0 - margin[0] - margin[2]) | |
| if cls.image_width is not None: | |
| kwargs['figsize'] = ( | |
| cls.image_width, | |
| cls.image_width * image_ratio / image_area_ratio | |
| ) | |
| elif cls.image_height: | |
| kwargs['figsize'] = ( | |
| cls.image_height * image_area_ratio / image_ratio, | |
| cls.image_height | |
| ) | |
| # dpi = max(cls.image_min_dpi, image.shape[1] / kwargs['figsize'][0] * cls.image_dpi_factor) | |
| dpi = 200 | |
| # import pdb; pdb.set_trace() | |
| fig = plt.figure(dpi=dpi, **kwargs) | |
| ax = plt.Axes(fig, [0.0 + margin[0], | |
| 0.0 + margin[1], | |
| 1.0 - margin[2], | |
| 1.0 - margin[3]]) | |
| ax.set_axis_off() | |
| ax.set_xlim(-0.5, image.shape[1] - 0.5) # imshow uses center-pixel-coordinates | |
| ax.set_ylim(image.shape[0] - 0.5, -0.5) | |
| fig.add_axes(ax) | |
| ax.imshow(image) | |
| if cls.white_overlay: | |
| white_screen(ax, cls.white_overlay) | |
| yield ax | |
| if fig_file: | |
| LOG.debug('writing image to %s', fig_file) | |
| fig.savefig(fig_file) | |
| if cls.show: | |
| plt.show() | |
| import pdb;pdb.set_trace() | |
| plt.close(fig) | |
| def white_screen(ax, alpha=0.9): | |
| ax.add_patch( | |
| plt.Rectangle((0, 0), 1, 1, transform=ax.transAxes, alpha=alpha, | |
| facecolor='white') | |
| ) | |
| canvas = Canvas.blank | |
| image_canvas = Canvas.image |