Source code for docarray.document.mixins.rich_embedding

import colorsys

from rich.color import Color
from rich.console import Console
from rich.console import ConsoleOptions, RenderResult
from rich.measure import Measurement
from rich.segment import Segment
from rich.style import Style

from docarray.math.helper import minmax_normalize
from docarray.math.ndarray import to_numpy_array


class ColorBoxEmbedding:
    def __init__(self, array):
        self._array = minmax_normalize(array, (0, 5))

    def __rich_console__(
        self, console: Console, options: ConsoleOptions
    ) -> RenderResult:
        h = 0.75
        for idx, y in enumerate(self._array):
            l = 0.1 + ((y / 5) * 0.7)
            r2, g2, b2 = colorsys.hls_to_rgb(h, l + 0.7 / 10, 1.0)
            color = Color.from_rgb(r2 * 255, g2 * 255, b2 * 255)
            yield Segment('▄', Style(color=color, bgcolor=color))
            if idx != 0 and idx % options.max_width == 0:
                yield Segment.line()

    def __rich_measure__(
        self, console: "Console", options: ConsoleOptions
    ) -> Measurement:
        return Measurement(1, options.max_width)