Source code for cutcutcodec.core.generation.video.matplotlib

"""Generate a video wih a matplotlib figure."""

from fractions import Fraction
import math
import threading
import typing

import matplotlib.figure
import matplotlib.pyplot
import torch

from cutcutcodec.core.classes.colorspace import Colorspace
from cutcutcodec.core.classes.container import ContainerInput
from cutcutcodec.core.classes.stream import Stream
from cutcutcodec.core.classes.stream_video import StreamVideo
from cutcutcodec.core.filter.video.resize import resize_keep_ratio




class GeneratorVideoMatplotlib(ContainerInput):
    """Generate a video from matplotlib figure.

    Examples
    --------
    >>> from fractions import Fraction
    >>> import matplotlib.pyplot as plt
    >>> from cutcutcodec.core.generation.video.matplotlib import GeneratorVideoMatplotlib
    >>>
    >>> def func(timestamp: Fraction) -> plt.Figure:
    ...     fig = plt.figure(layout="constrained", figsize=(7, 4))  # w, h
    ...     fig.supxlabel("abscissa")
    ...     fig.supylabel("ordinate")
    ...     axe = fig.subplots(squeeze=True)
    ...     axe.set_xlim([0, 1])
    ...     axe.set_ylim([0, 1])
    ...     axe.plot([0, 1], [timestamp, timestamp])  # horizontal line
    ...     return fig
    ...
    >>> (stream,) = GeneratorVideoMatplotlib(func).out_streams
    >>> stream.snapshot(0, (9, 9))
    >>>
    """

    def __init__(self, func: typing.Callable[[Fraction], matplotlib.figure.Figure]):
        """Initialise and create the class.

        Parameters
        ----------
        func : callable
            A user-defined function that takes in input the frame time as a rational number,
            and returns the filled matplotlib figure as the image of the current frame.

        """
        assert callable(func), func.__class__.__name__
        self.func = func
        super().__init__([_StreamVideoMatplotlib(self)])

    def _getstate(self) -> dict:
        return {"func": str(self.func)}

    def _setstate(self, in_streams: typing.Iterable[Stream], state: dict) -> None:
        assert state.keys() == {"func"}, set(state)
        self.func = state["func"]
        ContainerInput.__init__(self, [_StreamVideoNoiseUniform(self)])



class _StreamVideoMatplotlib(StreamVideo):
    """Matplotlib figure bases video stream."""

    colorspace = Colorspace.from_default_target_rgb()  # not working

    def __init__(self, node: GeneratorVideoMatplotlib):
        assert isinstance(node, GeneratorVideoMatplotlib), node.__class__.__name__
        super().__init__(node)

    def _snapshot(self, timestamp: Fraction, mask: torch.Tensor) -> torch.Tensor:
        if timestamp < 0:
            raise OutOfTimeRange(f"there is no video frame at timestamp {timestamp} (need >= 0)")
        if threading.current_thread().name != "MainThread":
            matplotlib.pyplot.switch_backend("agg")  # to allow matplotlib in thread
        fig = self.node.func(timestamp)
        # find the best dpi value to export the figure with the best resolution
        width_inches, height_inches = fig.get_size_inches()
        dpi = min(mask.shape[0]/height_inches, mask.shape[1]/width_inches)
        fig.set_dpi(dpi)
        # convert fig into frame
        fig.canvas.draw()
        width, height = fig.canvas.get_width_height()
        frame = (
            torch.frombuffer(bytearray(fig.canvas.tostring_argb()), dtype=torch.uint8)
            .reshape(height, width, 4)
            [:, :, 1:]  # argb to rgb
            .to(torch.float32)
        ) / 255.0
        matplotlib.pyplot.close(fig)  # fig.clear()
        return resize_keep_ratio(frame, mask.shape, copy=False)  # to guaranty having the write shape

    @property
    def beginning(self) -> Fraction:
        return Fraction(0)

    @property
    def duration(self) -> Fraction | float:
        return math.inf