ikvmocr

#!/usr/bin/env python3
# ikvmocr (part of ossobv/vcutil) // wdoekes/2022,2026 // Public Domain
#
# ikvmocr looks for a SuperMicro iKVM console window in the screenshot and does
# character recognition on it. It dumps the screenshot as characters to stdout.
#
# If you use GNOME, take a window-only screenshot of your iKVM window using
# Alt-PrintScreen and feed the saved PNG image to ikvmocr.
#
# An example ikvmocr.png is provided with this script to test things, as
# is a preloaded Glyph database, which may not be complete.
#
# Python dependencies:
# - Pillow (python3-pil)
#
# See also:
# - ipmikvm(1) to simply connecting to SuperMicro IPMI KVM from the console
# - xpaste(1) to paste characters to the Java iKVM program (kind of the
#   inverse of this, which reads characters from the program)
#
# Todo:
# - Add option to pass in glyph size and grid size manually.
# - Explain how glyphs are stored and add option to decode/show them.
# - Maybe crop glyphs before handling. (Requires looping over the entire
#   screenshot first. Or hardcoding glyph drop size in the DB.)
#
import json
import os
import sys
import time
from collections import defaultdict
from functools import wraps
from tempfile import NamedTemporaryFile

from PIL import Image

DATA_DIR = f'{os.environ.get("XDG_DATA_HOME", "~/.local/share")}/ikvmocr'


def timeit(title=None):
    def _wrapper(func):
        func_title = title or func.__name__

        @wraps(func)
        def _timer(*args, **kwargs):
            t0 = time.time()
            ret = func(*args, **kwargs)
            tn = time.time()
            print('timing {}: {:.3f} s'.format(
                func_title, tn - t0), file=sys.stderr)
            return ret
        return _timer
    return _wrapper


@timeit()
def detect_console_grid(pixels, width, height):
    # Calculate the sum of "lit" pixels for every single row and column.
    # We use at most 90% of the center of the window, so we skip borders.
    def norm(sum_, max_):
        return sum_ / 255 / max_

    xoff = width // 10
    smallwidth = width - 2 * xoff
    yoff = height // 10
    smallheight = height - 2 * yoff
    row_sums = [
        norm(sum(pixels[x + xoff, y + yoff]
                 for x in range(smallwidth)), smallwidth)
        for y in range(smallheight)]
    col_sums = [
        norm(sum(pixels[x + xoff, y + yoff]
                 for y in range(smallheight)), smallheight)
        for x in range(smallwidth)]
    # example: [..a bunch of floats, with some being (close to) zero..]

    # Check the max brightness on a line and pretend that all lines with
    # 20% of that are entirely dark.
    def binarize_lines(floats):
        pct20 = max(floats) / 5
        return [(0 if i < pct20 else 1) for i in floats]

    row_bins = binarize_lines(row_sums)
    col_bins = binarize_lines(col_sums)
    # example: [0, 0, 1, 1, 1, 1, 0, ...]

    # Start checking possible character size by taking a few appropriate
    # offsets and summing all consecutive lines the same distance.
    def calculate_stride_sums(bins, min=5, max=20):
        len_ = len(bins)
        strides = {}
        for off in range(5, max + 1):
            sums = []
            for start in range(off):
                sums.append(sum(bins[idx] for idx in range(start, len_, off)))
            strides[off] = sums

        return strides

    row_strides = calculate_stride_sums(row_bins, min=6, max=19)
    col_strides = calculate_stride_sums(col_bins, min=5, max=15)
    # example: {5: [33, 34, 31, 34, 33], ..., 8: [0, 0, 48, 4, 15, 27, 7, 40]}

    def get_best_strides(strides):
        def zero_percent(sums):
            return sum(not i for i in sums) / len(sums)

        # If we get multiple candidates, we get the smallest one by
        # doing ascending sort (and reversing the percentage).
        with_zero_percentages = [
            (1 - zero_percent(sums), idx) for idx, sums in strides.items()]
        with_zero_percentages.sort()

        # For now, just take the first one.
        # print(with_zero_percentages)
        return with_zero_percentages[0][1]

    row_size = get_best_strides(row_strides)
    col_size = get_best_strides(col_strides)
    # example: 8

    # Now that we have the character dimensions, we must go back and
    # find the most appropriate point to start on.
    row_stride = row_strides[row_size]
    col_stride = col_strides[col_size]
    # example: [0, 0, 48, 4, 15, 27, 7, 40]  # for "8"
    # (Ideally, we have exactly one 0, but no one has that kind of luck.)

    # For now, we'll center on the zeros in the middle. And then move
    # the char content to the bottom on a tie
    def find_stride_offset(stride):
        def get_max_consecutive_zeros(stride):
            best_start = max_zeros = current_zeros = 0
            # Loop through the list doubled to naturally handle the wrap-around
            for idx, num in enumerate(stride + stride):
                if current_zeros == 0:
                    current_start = idx
                if num == 0:
                    current_zeros += 1
                    if current_zeros > max_zeros:
                        max_zeros = current_zeros
                        best_start = current_start

                    max_zeros = max(current_zeros, max_zeros)
                else:
                    current_zeros = 0
            return best_start, max_zeros

        first_zero, count = get_max_consecutive_zeros(stride)
        center_offset = first_zero + (count // 2)  # prefer left zero
        # [17, 16, 17, 16, 16, 16, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 16, 16, 18]
        # has center_offset 11, so list[11:] + list[:11] is:
        # [0, 0, 0, 0, 0, 16, 16, 18, 17, 16, 17, 16, 16, 16, 17, 0, 0, 0, 0]
        return center_offset % len(stride)

    row_offset = find_stride_offset(row_stride)
    col_offset = find_stride_offset(col_stride)

    # Now we can finally go back and find the first valid line of the image.
    # Calculate offset direction first. Taking xoff and yoff into
    # account as well.
    validx = xoff + col_offset
    # Y detection seems off by one. Subtracting one gets us correct
    # values for all test images.
    validy = yoff + row_offset - 1

    # Search backwards and forwards.
    left = validx % col_size
    top = validy % row_size
    right = (width // col_size) * col_size + left
    if right > width:
        right -= col_size
    bottom = (height // row_size) * row_size + top
    if bottom > height:
        bottom -= row_size

    # Start at the first black line.
    while top < height:
        brightness = norm(sum(
            pixels[x + xoff, top] for x in range(smallwidth)), smallwidth)
        if brightness < 0.05:
            break
        top += row_size

    # Start at the first black line.
    while left < width:
        brightness = norm(sum(
            pixels[left, y + yoff] for y in range(smallheight)), smallheight)
        if brightness < 0.05:
            break
        left += col_size

    # End at the last black line.
    while bottom > top:
        bottom -= row_size
        brightness = norm(sum(
            pixels[x + xoff, bottom] for x in range(smallwidth)), smallwidth)
        if brightness < 0.05:
            break
    bottom += row_size

    # End at the last black line.
    while right > left:
        right -= col_size
        brightness = norm(sum(
            pixels[right, y + yoff] for y in range(smallheight)), smallheight)
        if brightness < 0.05:
            break
    right += col_size

    # print(
    #     (col_size, row_size), (validx, validy), col_stride, row_stride,
    #     (col_offset, row_offset))
    return (col_size, row_size), ((left, top), (right, bottom))


class ImageMixin:
    def __init__(self, pilimg):
        self._img = pilimg

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        self._img.close()
        self._img = None


class IKvmScreenshot(ImageMixin):
    """
    Take a "active window" screenshot using GNOME Alt-PrintScreen and
    supply the image to this instance.
    """
    @staticmethod
    def denoise_image(img, noise_threshold=48):
        """
        Return black/white in-memory image with noise cutoff.
        """
        return img.convert('L').point(
            (lambda p: 255 if p > noise_threshold else 0), mode='1')

    @classmethod
    def from_filename(cls, filename):
        return cls(Image.open(filename))

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Replace image and load pixels for quicker access.
        denoised_img = self.denoise_image(self._img)
        self._img.close()
        self._img = denoised_img

        # Get dimensions and pixels for quicker access.
        self._width, self._height = self._img.size
        self._pixels = self._img.load()

    def get_auto_console_grid(self):
        def to_gridsize(charsize, window, idx):
            assert (window[1][idx] - window[0][idx]) % charsize[idx] == 0, (
                charsize, window)
            return (window[1][idx] - window[0][idx]) // charsize[idx]

        charsize, window = detect_console_grid(
            self._pixels, self._width, self._height)
        gridsize = (
            to_gridsize(charsize, window, 0),
            to_gridsize(charsize, window, 1),
        )
        return self.get_console_grid(charsize, gridsize, window[0])

    def get_console_grid(self, charsize, gridsize, offset):
        return ConsoleGrid(self._pixels, charsize, gridsize, offset)


class ConsoleGrid:
    """
    Holds the console window with letters.
    """
    def __init__(self, pixels, charsize, gridsize, offset):
        self.pixels = pixels
        self.charsize = charsize
        self.gridsize = gridsize
        self.offset = offset

    def set_glyphs(self, glyphs):
        self.glyphs = glyphs
        self.glyphs.set_active_size(*self.charsize)

    @timeit()
    def get_content(self):
        lines = []
        for row in range(0, self.gridsize[1]):
            line = []
            for col in range(0, self.gridsize[0]):
                char = self.get_character(col, row)
                glyph_id = char.as_int63s()
                string = self.glyphs.get(glyph_id, as_string=char.as_string)
                line.append(string)
            lines.append(''.join(line))
        return '\n'.join(lines) + '\n'

    def get_character(self, col, row):
        pixels = []
        yoff = self.offset[1] + self.charsize[1] * row
        xoff = self.offset[0] + self.charsize[0] * col
        for y in range(yoff, yoff + self.charsize[1]):
            for x in range(xoff, xoff + self.charsize[0]):
                pixels.append(self.pixels[x, y])
        return ConsoleChar.from_pixels(self.charsize, pixels)


class ConsoleChar:
    """
    Holds a single glyph and does RLE encoding/decoding.
    """
    RLE_BITS = 4
    RLE_LTR = False
    BITS_PER_INT = 63

    # Don't touch.
    RLE_MASK = ((1 << RLE_BITS) - 1)
    BITS_PER_INT_MASK = ((1 << BITS_PER_INT) - 1)

    @classmethod
    def from_int63s(cls, colsize, rowsize, int63s):
        "Decode a list of 63 bit integers"
        # Use Python's big number capabilities and merge the integers
        bignumber = 0
        for idx, number in enumerate(int63s):
            assert number <= cls.BITS_PER_INT_MASK, number
            bignumber |= (number << (cls.BITS_PER_INT * idx))

        # Extract the N-bit chunks
        rle_bits = cls.RLE_BITS
        rle_mask = cls.RLE_MASK

        counts = []
        while bignumber:
            count = chunk = (bignumber & rle_mask)
            bignumber >>= rle_bits

            while chunk == rle_mask:
                has_more = (bignumber & 0x1)
                bignumber >>= 1
                if not has_more:
                    break
                chunk = (bignumber & rle_mask)
                assert chunk
                count += chunk
                bignumber >>= rle_bits
            assert not counts or chunk  # first one can be 0
            counts.append(count)

        # Turns counts into bool array
        bwdata = cls._from_counts(colsize, rowsize, counts)
        return ConsoleChar(colsize, rowsize, bwdata)

    @classmethod
    def from_pixels(cls, charsize, pixels):
        "Decode pixels into boolean on/off"
        colsize, rowsize = charsize
        bindata = []
        for pixel in pixels:
            bindata.append(bool(pixel))
        return cls(colsize, rowsize, bindata)

    def __init__(self, colsize, rowsize, bwdata):
        self.colsize = colsize
        self.rowsize = rowsize
        self._bwdata = bwdata

    def as_int63s(self):
        """Encode RLE counts as list of 63-bit integers, LSB-first.
        Each count is variable-length encoded: 5 data bits + 1 optional
        continuation bit, repeated as needed."""
        rle_bits = self.RLE_BITS
        rle_mask = self.RLE_MASK

        counts = self._to_counts(
            self.colsize, self.rowsize, self._bwdata)

        bitsum = bits = 0  # encode values per rle_bits bits
        for count in counts:
            assert not bits or count  # first one can be 0
            while count:
                chunk = min(count, rle_mask)
                count -= chunk
                bitsum |= (chunk << bits)
                bits += rle_bits
                if chunk == rle_mask:
                    if count > 0:
                        bitsum |= (0x1 << bits)
                    bits += 1

        int63s = []
        while bitsum:
            int63s.append(bitsum & self.BITS_PER_INT_MASK)
            bitsum >>= self.BITS_PER_INT

        return tuple(int63s)

    def as_string(self):
        line = []
        lines = []
        for idx, cur in enumerate(self._bwdata, 1):
            line.append('[X]' if cur else ' - ')
            if idx % self.colsize == 0:
                line.append('|')
                lines.append(''.join(line))
                line = []
        assert line == [], line
        return '\n'.join(lines)

    @classmethod
    def _from_counts(cls, colsize, rowsize, counts):
        "Decode counts to list of true/false values, first false"
        if cls.RLE_LTR:
            return cls._from_counts_horizontal(colsize, rowsize, counts)
        return cls._from_counts_vertical(colsize, rowsize, counts)

    @staticmethod
    def _from_counts_horizontal(colsize, rowsize, counts):
        "Decode counts to list of true/false values, first false, LTR-TTB"
        total_length = colsize * rowsize

        bwdata = []
        state = False
        for count in counts:
            bwdata.extend([state] * count)
            state = not state

        # Fill the remaining length with the final state
        bwdata.extend([state] * (total_length - len(bwdata)))
        return bwdata

    @staticmethod
    def _from_counts_vertical(colsize, rowsize, counts):
        "Decode counts to list of true/false values, first false, TTB-LTR"
        total_length = colsize * rowsize
        bwdata = [False] * total_length
        state = False
        pos = 0
        for count in counts:
            if state:
                for _ in range(count):
                    col = pos // rowsize
                    row = pos % rowsize
                    bwdata[row * colsize + col] = True
                    pos += 1
            else:
                pos += count
            state = not state
        if state:
            while pos < total_length:
                col = pos // rowsize
                row = pos % rowsize
                bwdata[row * colsize + col] = True
                pos += 1
        return bwdata

    @classmethod
    def _to_counts(cls, colsize, rowsize, bwdata):
        "Encode as counts of true/false, starting with false"
        if cls.RLE_LTR:
            return cls._to_counts_horizontal(colsize, rowsize, bwdata)
        return cls._to_counts_vertical(colsize, rowsize, bwdata)

    @staticmethod
    def _to_counts_horizontal(colsize, rowsize, bwdata):
        "Encode as counts of true/false, starting with false, LTR-TTB."
        ret = []
        last, count = False, 0
        for cur in bwdata:
            if cur == last:
                count += 1
            else:
                ret.append(count)
                last, count = cur, 1
        # NOTE: We purposefully do NOT add the last count. We don't need
        # it. (It follows from the previous values along with the
        # character dimensions.)
        return ret

    @staticmethod
    def _to_counts_vertical(colsize, rowsize, bwdata):
        "Encode as counts of true/false, starting with false, TTB-LTR."
        ret = []
        last, count = False, 0
        for col in range(colsize):
            for row in range(rowsize):
                cur = bwdata[row * colsize + col]
                if cur == last:
                    count += 1
                else:
                    ret.append(count)
                    last, count = cur, 1
        # NOTE: We purposefully do NOT add the last count. We don't need
        # it. (It follows from the previous values along with the
        # character dimensions.)
        return ret


class ConsoleGlyphs:
    """
    Holds the set of glyphs for all different character sizes and loads/saves.
    """
    def __init__(self):
        self._storage = {}
        self._active_cols_rows = None

    @timeit('(the loading of glyphs from disk)')
    def load_from_file(self, fp):
        data = json.load(fp)
        storage = {}
        for cols_rows, js_array in data.items():
            cols, rows = [int(i) for i in cols_rows.split(',')]
            storage[(cols, rows)] = self._alpha_js_to_mem(js_array)
        self._storage = storage

    def save_to_file(self, fp):
        """This does manual JS writing so we get a pretty/readable DB.
        We wanted to store the int63s as 0xNUMBER, but JS does not allow
        0x notation, so now it's a hex string instead (see
        _str_from_glyph_id)."""
        fp.write('{\n')
        for idx1, ((cols, rows), mem_dict) in enumerate(
                sorted(self._storage.items()), 1):
            trailing_comma1 = ',' if idx1 != len(self._storage) else ''
            fp.write('  "{},{}": {{\n'.format(cols, rows))
            js_dict = self._alpha_mem_to_js(mem_dict)
            for idx2, (string, glyph_ids) in enumerate(
                    sorted(js_dict.items()), 1):
                glyph_format = '", "'.join(sorted(glyph_ids))
                trailing_comma2 = ',' if idx2 != len(js_dict) else ''
                fp.write('    {}: ["{}"]{}\n'.format(
                    json.dumps(string), glyph_format, trailing_comma2))
            fp.write('  }}{}\n'.format(trailing_comma1))
        fp.write('}\n')

    def set_active_size(self, colsize, rowsize):
        self._active_cols_rows = (colsize, rowsize)
        if self._active_cols_rows not in self._storage:
            self._storage[self._active_cols_rows] = {}

    def get(self, glyph_id, as_string=None):
        try:
            ch = self._storage[self._active_cols_rows][glyph_id]
        except KeyError:
            assert callable(as_string), (glyph_id, as_string)
            print('UNKNOWN GLYPH_ID:', self._str_from_glyph_id(glyph_id))
            print('UNKNOWN GLYPH LAYOUT:')
            print(as_string())
            if os.environ.get('RUNTESTS', '') in ('', '0'):
                ch = input('Specify which character: ')
                if ch.startswith(' '):
                    ch = ' '
                else:
                    ch = ch.strip()
                assert ch != '', repr(ch)
            else:
                # In test suite, use fail-char
                ch = '\ufffd'  # Unicode fail question mark
            self.set(glyph_id, ch)
        return ch

    def set(self, glyph_id, string):
        assert glyph_id not in self._storage[self._active_cols_rows]
        self._storage[self._active_cols_rows][glyph_id] = string

    def _str_from_glyph_id(self, glyph_id):
        "[0x123, 0x456] => '123:456'"
        glyph_parts = ['{:x}'.format(gl) for gl in glyph_id]
        return ':'.join(glyph_parts)

    def _str_to_glyph_id(self, glyph_id):
        "'123:456' => [0x123, 0x456]"
        if not glyph_id:
            return ()
        return tuple(int(gl, 16) for gl in glyph_id.split(':'))

    def _alpha_mem_to_js(self, mem_dict):
        ret = defaultdict(list)
        for glyph_id, string in mem_dict.items():
            ret[string].append(self._str_from_glyph_id(glyph_id))
        return ret

    def _alpha_js_to_mem(self, js_dict):
        ret = {}
        for string, str_glyph_ids in js_dict.items():
            for str_glyph_id in str_glyph_ids:
                glyph_id = self._str_to_glyph_id(str_glyph_id)
                assert glyph_id not in ret, (string, str_glyph_id, glyph_id)
                ret[glyph_id] = string
        return ret


def truncate_screenshot(text):
    lines = text.split('\n')

    # Remove leading/trailing all-whitespace lines
    while lines and not lines[0].strip():
        lines.pop(0)
    while lines and not lines[-1].strip():
        lines.pop()

    if not lines:
        return ''

    # Find common indent from non-blank lines only.
    indent = min(
        (len(line) - len(line.lstrip(' ')))
        for line in lines if line.strip())
    outdent = max(len(line.rstrip(' ')) for line in lines if line.strip())

    lines = [line[indent:outdent] for line in lines]
    return '\n'.join(lines) + '\n'


def read_screenshot(screenshot_filename, glyphs):
    with IKvmScreenshot.from_filename(screenshot_filename) as screenshot:
        consolegrid = screenshot.get_auto_console_grid()
        consolegrid.set_glyphs(glyphs)
        return truncate_screenshot(consolegrid.get_content())


def read_screenshot_and_dump(screenshot_filename, glyphs):
    content = read_screenshot(screenshot_filename, glyphs)
    print('----', screenshot_filename, '----', file=sys.stderr)
    print(content, end='')


def main(screenshot_filenames, glyph_db_filename):
    glyph_db_filenames = [
        glyph_db_filename,
        os.path.expanduser(f'{DATA_DIR}/ikvmocr.js'),
    ]
    glyphs = ConsoleGlyphs()
    for glyph_db_filename in glyph_db_filenames:
        try:
            with open(glyph_db_filename, 'r') as fp:
                glyphs.load_from_file(fp)
        except FileNotFoundError:
            pass
        else:
            break
    else:
        print('no glyph db found at {}, CREATING NEW'.format(
            glyph_db_filename), file=sys.stderr)
        try:
            os.makedirs(os.path.dirname(glyph_db_filename))
        except FileExistsError:
            pass

    if 0:
        glyphs.set_active_size(8, 16)
        g = glyphs.get((0x718e3ce3cf3cf3cf, 0x279e09))
        print(g)
        q = ConsoleChar.from_int63s(
            8, 16, glyphs._str_to_glyph_id('718e3ce3cf3cf3cf:279e09'))
        print(q)
        print(q.as_string())

    try:
        for screenshot_filename in screenshot_filenames:
            read_screenshot_and_dump(screenshot_filename, glyphs)
    finally:
        with NamedTemporaryFile(
                prefix=(os.path.basename(glyph_db_filename) + '.'),
                dir=os.path.dirname(glyph_db_filename),
                mode='w', delete=False) as fp:
            try:
                glyphs.save_to_file(fp)
            except BaseException:
                os.unlink(fp.name)
                raise
        os.rename(fp.name, glyph_db_filename)


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print(f'''\
usage: ikvmocr SCREENSHOT_FILE...

ikvmocr looks for a SuperMicro iKVM console window in the screenshot and does
character recognition on it. It dumps the screenshot as characters to stdout.

If you use GNOME, take a window-only screenshot of your iKVM window using
Alt-PrintScreen and feed the saved PNG image to ikvmocr.

The glyph configuration file is taken from the same directory as the binary,
if it exists. Otherwise it uses {DATA_DIR}/ikvmocr.js.

See also:
- ipmikvm(1) to simply connecting to SuperMicro IPMI KVM from the console
- xpaste(1) to paste characters to the Java iKVM program (the inverse of this)
''', file=sys.stderr)
        sys.exit(1)

    glyph_db_filename = os.path.join(os.path.dirname(__file__), 'ikvmocr.js')
    main(sys.argv[1:], glyph_db_filename)