Source code for kornia.utils.image_print

"""Convert an image tensor to an ANSI text string (xterm-256color).

Nice long listing of all 256 colors and their codes.

Taken from https://gist.github.com/klange/1687427
"""
import re
from typing import Tuple, Union

from torch import float16, float32, float64

import kornia
from kornia.core import Tensor
from kornia.core.check import KORNIA_CHECK_IS_IMAGE, KORNIA_CHECK_SHAPE
from kornia.io import ImageLoadType

# color look-up table
# 8-bit, RGB hex
CLUT = [
    # Primary 3-bit (8 colors). Unique representation!
    ("00", "000000"),
    ("01", "800000"),
    ("02", "008000"),
    ("03", "808000"),
    ("04", "000080"),
    ("05", "800080"),
    ("06", "008080"),
    ("07", "c0c0c0"),
    # Equivalent "bright" versions of original 8 colors.
    ("08", "808080"),
    ("09", "ff0000"),
    ("10", "00ff00"),
    ("11", "ffff00"),
    ("12", "0000ff"),
    ("13", "ff00ff"),
    ("14", "00ffff"),
    ("15", "ffffff"),
    # Strictly ascending.
    ("16", "000000"),
    ("17", "00005f"),
    ("18", "000087"),
    ("19", "0000af"),
    ("20", "0000d7"),
    ("21", "0000ff"),
    ("22", "005f00"),
    ("23", "005f5f"),
    ("24", "005f87"),
    ("25", "005faf"),
    ("26", "005fd7"),
    ("27", "005fff"),
    ("28", "008700"),
    ("29", "00875f"),
    ("30", "008787"),
    ("31", "0087af"),
    ("32", "0087d7"),
    ("33", "0087ff"),
    ("34", "00af00"),
    ("35", "00af5f"),
    ("36", "00af87"),
    ("37", "00afaf"),
    ("38", "00afd7"),
    ("39", "00afff"),
    ("40", "00d700"),
    ("41", "00d75f"),
    ("42", "00d787"),
    ("43", "00d7af"),
    ("44", "00d7d7"),
    ("45", "00d7ff"),
    ("46", "00ff00"),
    ("47", "00ff5f"),
    ("48", "00ff87"),
    ("49", "00ffaf"),
    ("50", "00ffd7"),
    ("51", "00ffff"),
    ("52", "5f0000"),
    ("53", "5f005f"),
    ("54", "5f0087"),
    ("55", "5f00af"),
    ("56", "5f00d7"),
    ("57", "5f00ff"),
    ("58", "5f5f00"),
    ("59", "5f5f5f"),
    ("60", "5f5f87"),
    ("61", "5f5faf"),
    ("62", "5f5fd7"),
    ("63", "5f5fff"),
    ("64", "5f8700"),
    ("65", "5f875f"),
    ("66", "5f8787"),
    ("67", "5f87af"),
    ("68", "5f87d7"),
    ("69", "5f87ff"),
    ("70", "5faf00"),
    ("71", "5faf5f"),
    ("72", "5faf87"),
    ("73", "5fafaf"),
    ("74", "5fafd7"),
    ("75", "5fafff"),
    ("76", "5fd700"),
    ("77", "5fd75f"),
    ("78", "5fd787"),
    ("79", "5fd7af"),
    ("80", "5fd7d7"),
    ("81", "5fd7ff"),
    ("82", "5fff00"),
    ("83", "5fff5f"),
    ("84", "5fff87"),
    ("85", "5fffaf"),
    ("86", "5fffd7"),
    ("87", "5fffff"),
    ("88", "870000"),
    ("89", "87005f"),
    ("90", "870087"),
    ("91", "8700af"),
    ("92", "8700d7"),
    ("93", "8700ff"),
    ("94", "875f00"),
    ("95", "875f5f"),
    ("96", "875f87"),
    ("97", "875faf"),
    ("98", "875fd7"),
    ("99", "875fff"),
    ("100", "878700"),
    ("101", "87875f"),
    ("102", "878787"),
    ("103", "8787af"),
    ("104", "8787d7"),
    ("105", "8787ff"),
    ("106", "87af00"),
    ("107", "87af5f"),
    ("108", "87af87"),
    ("109", "87afaf"),
    ("110", "87afd7"),
    ("111", "87afff"),
    ("112", "87d700"),
    ("113", "87d75f"),
    ("114", "87d787"),
    ("115", "87d7af"),
    ("116", "87d7d7"),
    ("117", "87d7ff"),
    ("118", "87ff00"),
    ("119", "87ff5f"),
    ("120", "87ff87"),
    ("121", "87ffaf"),
    ("122", "87ffd7"),
    ("123", "87ffff"),
    ("124", "af0000"),
    ("125", "af005f"),
    ("126", "af0087"),
    ("127", "af00af"),
    ("128", "af00d7"),
    ("129", "af00ff"),
    ("130", "af5f00"),
    ("131", "af5f5f"),
    ("132", "af5f87"),
    ("133", "af5faf"),
    ("134", "af5fd7"),
    ("135", "af5fff"),
    ("136", "af8700"),
    ("137", "af875f"),
    ("138", "af8787"),
    ("139", "af87af"),
    ("140", "af87d7"),
    ("141", "af87ff"),
    ("142", "afaf00"),
    ("143", "afaf5f"),
    ("144", "afaf87"),
    ("145", "afafaf"),
    ("146", "afafd7"),
    ("147", "afafff"),
    ("148", "afd700"),
    ("149", "afd75f"),
    ("150", "afd787"),
    ("151", "afd7af"),
    ("152", "afd7d7"),
    ("153", "afd7ff"),
    ("154", "afff00"),
    ("155", "afff5f"),
    ("156", "afff87"),
    ("157", "afffaf"),
    ("158", "afffd7"),
    ("159", "afffff"),
    ("160", "d70000"),
    ("161", "d7005f"),
    ("162", "d70087"),
    ("163", "d700af"),
    ("164", "d700d7"),
    ("165", "d700ff"),
    ("166", "d75f00"),
    ("167", "d75f5f"),
    ("168", "d75f87"),
    ("169", "d75faf"),
    ("170", "d75fd7"),
    ("171", "d75fff"),
    ("172", "d78700"),
    ("173", "d7875f"),
    ("174", "d78787"),
    ("175", "d787af"),
    ("176", "d787d7"),
    ("177", "d787ff"),
    ("178", "d7af00"),
    ("179", "d7af5f"),
    ("180", "d7af87"),
    ("181", "d7afaf"),
    ("182", "d7afd7"),
    ("183", "d7afff"),
    ("184", "d7d700"),
    ("185", "d7d75f"),
    ("186", "d7d787"),
    ("187", "d7d7af"),
    ("188", "d7d7d7"),
    ("189", "d7d7ff"),
    ("190", "d7ff00"),
    ("191", "d7ff5f"),
    ("192", "d7ff87"),
    ("193", "d7ffaf"),
    ("194", "d7ffd7"),
    ("195", "d7ffff"),
    ("196", "ff0000"),
    ("197", "ff005f"),
    ("198", "ff0087"),
    ("199", "ff00af"),
    ("200", "ff00d7"),
    ("201", "ff00ff"),
    ("202", "ff5f00"),
    ("203", "ff5f5f"),
    ("204", "ff5f87"),
    ("205", "ff5faf"),
    ("206", "ff5fd7"),
    ("207", "ff5fff"),
    ("208", "ff8700"),
    ("209", "ff875f"),
    ("210", "ff8787"),
    ("211", "ff87af"),
    ("212", "ff87d7"),
    ("213", "ff87ff"),
    ("214", "ffaf00"),
    ("215", "ffaf5f"),
    ("216", "ffaf87"),
    ("217", "ffafaf"),
    ("218", "ffafd7"),
    ("219", "ffafff"),
    ("220", "ffd700"),
    ("221", "ffd75f"),
    ("222", "ffd787"),
    ("223", "ffd7af"),
    ("224", "ffd7d7"),
    ("225", "ffd7ff"),
    ("226", "ffff00"),
    ("227", "ffff5f"),
    ("228", "ffff87"),
    ("229", "ffffaf"),
    ("230", "ffffd7"),
    ("231", "ffffff"),
    # Gray-scale range.
    ("232", "080808"),
    ("233", "121212"),
    ("234", "1c1c1c"),
    ("235", "262626"),
    ("236", "303030"),
    ("237", "3a3a3a"),
    ("238", "444444"),
    ("239", "4e4e4e"),
    ("240", "585858"),
    ("241", "626262"),
    ("242", "6c6c6c"),
    ("243", "767676"),
    ("244", "808080"),
    ("245", "8a8a8a"),
    ("246", "949494"),
    ("247", "9e9e9e"),
    ("248", "a8a8a8"),
    ("249", "b2b2b2"),
    ("250", "bcbcbc"),
    ("251", "c6c6c6"),
    ("252", "d0d0d0"),
    ("253", "dadada"),
    ("254", "e4e4e4"),
    ("255", "eeeeee"),
]
SHORT2RGB_DICT = dict(CLUT)
RGB2SHORT_DICT = {v: k for k, v in SHORT2RGB_DICT.items()}


def _str2hex(hexstr: str) -> int:
    return int(hexstr, 16)


def _strip_hash(rgb: str) -> str:
    # Strip leading `#` if exists.
    if rgb.startswith("#"):
        rgb = rgb.lstrip("#")
    return rgb


def short2rgb(short: str) -> str:
    return SHORT2RGB_DICT[short]


def rgb2short(rgb: str) -> Tuple[str, str]:
    """Find the closest xterm-256 approximation to the given RGB value.

    Args:
        rgb: Hex code representing an RGB value, eg, 'abcdef'.

    Returns:
        String between 0 and 255, compatible with xterm.

    Example:
        >>> rgb2short('123456')
        ('23', '005f5f')
        >>> rgb2short('ffffff')
        ('231', 'ffffff')
        >>> rgb2short('0DADD6')  # vimeo logo
        ('38', '00afd7')
    """
    rgb = _strip_hash(rgb)
    incs = (0x00, 0x5F, 0x87, 0xAF, 0xD7, 0xFF)
    # Break 6-char RGB code into 3 integer vals.
    parts = [int(h, 16) for h in re.split(r"(..)(..)(..)", rgb)[1:4]]
    res = []
    for part in parts:
        i = 0
        while i < len(incs) - 1:
            s, b = incs[i], incs[i + 1]  # smaller, bigger
            if s <= part <= b:
                s1 = abs(s - part)
                b1 = abs(b - part)
                if s1 < b1:
                    closest = s
                else:
                    closest = b
                res.append(closest)
                break
            i += 1
    _res = "".join([("%02.x" % i) for i in res])
    equiv = RGB2SHORT_DICT[_res]
    return equiv, _res


[docs]def image_to_string(image: Tensor, max_width: int = 256) -> str: """Obtain the closest xterm-256 approximation string from an image tensor. The tensor shall be either 0~1 float type or 0~255 long type. Args: image: an RGB image with shape :math:`3HW`. max_width: maximum width of the input image. """ KORNIA_CHECK_IS_IMAGE(image, None, raises=True) KORNIA_CHECK_SHAPE(image, ["C", "H", "W"]) if image.dtype not in [float16, float32, float64]: image = image / 255.0 # In case of resizing. if image.shape[-1] > max_width: image = kornia.geometry.resize(image, (image.size(-2) * max_width // image.size(-1), max_width)) image = (image * 255).long() res = "" for y in range(image.size(-2)): for x in range(image.size(-1)): r, g, b = image[:, y, x] h = f"{r:2x}{g:2x}{b:2x}" short, _ = rgb2short(h) res += "\033[48;5;%sm " % short res += "\033[0m\n" return res