Source code for slitflow.fig.figure

import io
import os
import tempfile

import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt

from ..img.image import RGB
from ..data import Data, Pickle
from .. import MATPLOTLIB_BACKEND

matplotlib.use(MATPLOTLIB_BACKEND)




class Figure(Pickle):
    """Figure Data saved as matplotlib pickle object.

    """
    EXT = '.fig'

    def __init__(self, info_path=None):
        super().__init__(info_path)



    def save_data(self, data, path):
        """
        :class:`matplotlib.axes.Axes` cannot be copied. All figures should be
        saved as a file after executing each process. You have to use run
        mode 2.
        """
        super().save_data(data, path)
        plt.close()






class ToTiff(RGB):
    """Convert matplotlib figure object to RGB tiff class.

    Args:
        reqs[0] (Figure): Figure class containing
            :class:`matplotlib.figure.Figure` objects.
        param["dpi"] (int, optional): Dot per inch of image. Defaults to 400.
        param["scalebar"] (list, optional): [length, left position, bottom
            position, line width, line color] of the scale bar. Positions
            should be the relative positions of figure size (0-1). Line color
            should be [R(0-255), G(0-255), B(0-255)]. This parameter requires
            "limit", "size" and "length_unit" to reqs[0].
        param["split_depth"] (int): File split depth number.

    Returns:
        slitflow.img.image.RGB: Image of figure
    """



    def set_info(self, param={}):
        """Copy information from reqs[0] and add parameters.
        """
        if self.reqs[0].info.data_split_depth != \
                self.reqs[0].info.split_depth():
            raise Exception("Data should split to the file split depth")
        self.info.copy_req(0)
        if "dpi" not in param:
            param["dpi"] = 400
        self.info.add_param(
            "dpi", param["dpi"], "int", "Dot per inch of figure image")

        # use only x-axis for scale bar
        limit = self.info.get_param_value("limit")
        size = self.info.get_param_value("size")
        length_unit = self.info.get_param_value("length_unit")
        if limit and size and length_unit:
            pixel = size[0] / 2.54 * param["dpi"]
            if limit[0] is not None:
                pitch = (limit[1] - limit[0]) / pixel
                self.info.add_param(
                    "pitch", pitch, length_unit, "Pixel size of figure")
        if "scalebar" in param:
            self.info.add_param(
                "scalebar", param["scalebar"], "list",
                "[length, left, bottom , line_width, line_color] of the scale\
                 bar")
        self.info.set_split_depth(param["split_depth"])




    @staticmethod
    def process(reqs, param):
        """Convert matplotlib figure object to RGB tiff class.

        Args:
            reqs[0] (matplotlib.figure.Figure): Figure object.
            param["dpi"] (int, optional): Dot per inch of image. Defaults to
                400.
            param["scalebar"] (list, optional): [length, left position, bottom
                position, line width, line color] of the scale bar. Positions
                should be the relative positions of figure size (0-1).
                Line color should be [R(0-255), G(0-255), B(0-255)]. This
                parameter requires "limit", "size" and "length_unit" to
                reqs[0].
            param["limit"] (list of float, optional): [left, right, bottom,
                top] limits of figure axes. Required if ``scale`` in param.

        Returns:
            numpy.ndarray: Image of figure
        """
        fig = reqs[0]
        if "scalebar" in param:
            width = param["limit"][1] - param["limit"][0]
            height = param["limit"][3] - param["limit"][2]
            left = width * param["scalebar"][1] + param["limit"][0]
            bottom = height * param["scalebar"][2] + param["limit"][2]

            fig.axes[0].plot([left, left + param["scalebar"][0]],
                             [bottom, bottom], linewidth=param["scalebar"][3],
                             color=np.array(param["scalebar"][4]) / 255)

        stack = np.array(get_stack(fig, param["dpi"]))
        return stack






def get_stack(fig, dpi):
    """Convert matplotlib.figure.Figure to numpy.ndarray of RGB image.

    Args:
        fig (matplotlib.figure.Figure): Figure object.
        dpi (int): Dot per inch.

    Returns:
        numpy.ndarray: RGB image with the shape of (color, height, width)
    """
    stack = []
    fig.set_dpi(dpi)
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi)
    buf.seek(0)
    fig_rgb = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    img = cv2.imdecode(fig_rgb, 1)
    fig_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    stack.append(np.flipud(fig_rgb[:, :, 0]))
    stack.append(np.flipud(fig_rgb[:, :, 1]))
    stack.append(np.flipud(fig_rgb[:, :, 2]))
    return stack





def inherit_split_depth(Data, reqs_no, group_depth):
    """Set the group and split depth based on the reqs data.

    Args:
        Data (Data): Target Data object to set split_depth.
        reqs_no (int): Required Data number to get split_depth.
        group_depth (int): Data grouping depth number to set calc_cols.
    """
    Data.info.set_group_depth(group_depth)
    split_depth = Data.reqs[reqs_no].info.data_split_depth
    keeps = Data.info.get_column_name("index")[:split_depth]
    Data.info.delete_column(keeps=keeps)
    Data.info.set_split_depth(split_depth)