Source code for blocks.utils.demo

from typing import Sequence, Union

import numpy as np

from .base64 import cv2_to_base64


__all__ = ['DemoOutputBuilder']


class DemoOutputBuilder:
    """Building demo outputs to be displayed in Iterait `demo web application <https://github.com/Iterait/demotool>`_.


    .. code-block:: python
        :caption: example usage

        output = blocks.utils.DemoOutputBuilder().add_image(im).add_table(['col1', 'col2'], [['lorem ipsum', im2]]).output
        # output ~ demo-compatible json-serializable dict

    """

    def __init__(self):
        self._blocks = []

    def add_image(self, image: np.ndarray, extension: str='.jpg') -> 'DemoOutputBuilder':
        """Add image block to the output."""
        self._blocks.append(
            {
                "type": "image",
                "contentType": "image/png",
                "content": cv2_to_base64(image, extension=extension)
            })
        return self

    def add_image_transformation(self, original: np.ndarray, transformed: np.ndarray, extension: str='.jpg')\
            -> 'DemoOutputBuilder':
        """Add image comparison block to the output."""
        self._blocks.append(
            {
                "type": "imageTransformation",
                "original": {
                    "contentType": "image/png",
                    "content": cv2_to_base64(original, extension=extension)
                },
                "result": {
                    "contentType": "image/png",
                    "content": cv2_to_base64(transformed, extension=extension)
                }
            })
        return self

    def add_table(self, header: Sequence[str], rows: Sequence[Sequence[Union[str, np.ndarray]]]) \
            -> 'DemoOutputBuilder':
        """
        Add table block to the output.

        :param header: table header - a sequence of strings
        :param rows: a sequence of rows wherein row is a sequence of strings/images
        """
        output_rows = []
        for row in rows:
            output_row = []
            for cell in row:
                if isinstance(cell, str):
                    output_row.append(
                        {
                            "type": "text",
                            "content": cell
                        })
                elif isinstance(cell, np.ndarray):
                    output_row.append(
                        {
                            "type": "image",
                            "contentType": "image/png",
                            "content": cv2_to_base64(cell, extension='.png')
                        })
                else:
                    raise ValueError('Value type `{}` is not allowed in the demo table '
                                     '(only string and images are supported).'.format(type(cell)))
            output_rows.append(output_row)

        self._blocks.append(
            {
                "type": "table",
                "header": header,
                "rows": output_rows
            })
        return self

    @property
    def output(self) -> dict:
        """
        Return demo-compatible json-serializable dict.

        @teyras suggests to use this in dataset's `postprocess_batch`, **shepherd** runner or **emloop** model.
        """
        return {'contents': [{'blocks': self._blocks}]}