image_processing.py

"""TensorFlow-free image preprocessing and symbol segmentation."""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass

import numpy as np
from PIL import Image

IMG_SIZE = 28
GLYPH_SIZE = 20
PREPROCESSING_VERSION = 2


@dataclass(frozen=True)
class Component:
    x0: int
    y0: int
    x1: int
    y1: int
    area: int

    @property
    def width(self) -> int:
        return self.x1 - self.x0

    @property
    def height(self) -> int:
        return self.y1 - self.y0


def _ink_mask(pil: Image.Image) -> tuple[np.ndarray, bool]:
    arr = np.asarray(pil.convert("L"), dtype=np.uint8)
    if arr.size == 0:
        return np.zeros(arr.shape, dtype=bool), True
    border = np.concatenate((arr[0], arr[-1], arr[:, 0], arr[:, -1]))
    background = float(np.median(border))
    histogram = np.bincount(arr.ravel(), minlength=256)
    global_mode = int(np.argmax(histogram))
    # A glyph can touch or even surround every border pixel. In that case the
    # dominant interior color is a better background estimate than the border.
    if abs(global_mode - background) > 64 and histogram[global_mode] / arr.size >= 0.35:
        background = float(global_mode)
    light_background = background >= 128
    threshold = max(20, background - 35) if light_background else min(235, background + 35)
    mask = arr < threshold if light_background else arr > threshold
    if not mask.any():
        light_background = float(arr.mean()) > 127
        mask = arr < 245 if light_background else arr > 10
    return mask, light_background


def pil_to_arr(pil: Image.Image) -> np.ndarray:
    """Normalize a drawing into bright ink centered on a 28×28 black canvas."""
    arr = np.asarray(pil.convert("L"), dtype=np.uint8)
    if arr.size == 0:
        return np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.float32)

    mask, light_background = _ink_mask(pil)
    if not mask.any():
        return np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.float32)

    ys, xs = np.where(mask)
    cropped = arr[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1]
    if light_background:
        cropped = (255 - cropped).astype(np.uint8, copy=False)
    h, w = cropped.shape
    if h <= 0 or w <= 0:
        return np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.float32)

    size = max(h, w)
    square = np.zeros((size, size), dtype=np.uint8)
    y_pad, x_pad = (size - h) // 2, (size - w) // 2
    square[y_pad : y_pad + h, x_pad : x_pad + w] = cropped

    glyph = Image.fromarray(square).resize((GLYPH_SIZE, GLYPH_SIZE), Image.Resampling.LANCZOS)
    canvas = Image.new("L", (IMG_SIZE, IMG_SIZE), 0)
    offset = (IMG_SIZE - GLYPH_SIZE) // 2
    canvas.paste(glyph, (offset, offset))
    return np.asarray(canvas, dtype=np.float32) / 255.0


def find_strokes_bbox(pil: Image.Image) -> tuple[int, int, int, int] | None:
    mask, _ = _ink_mask(pil)
    if not mask.any():
        return None
    ys, xs = np.where(mask)
    # PIL crop coordinates are exclusive at the right and bottom.
    return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)


def connected_components(mask: np.ndarray, min_area: int = 2) -> list[Component]:
    """Return 8-connected ink components in deterministic left-to-right order."""
    if mask.ndim != 2:
        raise ValueError("mask must be a 2-D array")
    height, width = mask.shape
    seen = np.zeros_like(mask, dtype=bool)
    components: list[Component] = []

    for start_y, start_x in zip(*np.where(mask), strict=True):
        if seen[start_y, start_x]:
            continue
        queue = deque([(int(start_y), int(start_x))])
        seen[start_y, start_x] = True
        x0 = x1 = int(start_x)
        y0 = y1 = int(start_y)
        area = 0
        while queue:
            y, x = queue.pop()
            area += 1
            x0, x1 = min(x0, x), max(x1, x)
            y0, y1 = min(y0, y), max(y1, y)
            for next_y in range(max(0, y - 1), min(height, y + 2)):
                for next_x in range(max(0, x - 1), min(width, x + 2)):
                    if mask[next_y, next_x] and not seen[next_y, next_x]:
                        seen[next_y, next_x] = True
                        queue.append((next_y, next_x))
        if area >= min_area:
            components.append(Component(x0, y0, x1 + 1, y1 + 1, area))

    return sorted(components, key=lambda component: (component.x0, component.y0))


def _should_merge(previous: Component, current: Component, adaptive_gap: int) -> bool:
    horizontal_gap = current.x0 - previous.x1
    overlaps_x = current.x0 <= previous.x1
    vertical_overlap = min(previous.y1, current.y1) - max(previous.y0, current.y0)
    min_height = max(1, min(previous.height, current.height))
    overlaps_y_enough = vertical_overlap / min_height >= 0.20

    # Vertically separated pieces with x overlap form glyphs such as =, i, or ÷.
    if overlaps_x:
        return True
    # Horizontally nearby pieces are merged only when their vertical spans agree.
    return horizontal_gap <= adaptive_gap and overlaps_y_enough


def segment_symbols(pil: Image.Image, gap_threshold: int = 20) -> list[tuple[int, Image.Image]]:
    """Segment a line of handwriting using adaptive connected-component groups."""
    if gap_threshold < 0:
        raise ValueError("gap_threshold must be non-negative")
    mask, _ = _ink_mask(pil)
    components = connected_components(mask)
    if not components:
        return []

    median_width = float(np.median([component.width for component in components]))
    adaptive_gap = max(2, min(gap_threshold, round(median_width * 0.35)))
    groups: list[Component] = []
    for component in components:
        if not groups or not _should_merge(groups[-1], component, adaptive_gap):
            groups.append(component)
            continue
        previous = groups[-1]
        groups[-1] = Component(
            min(previous.x0, component.x0),
            min(previous.y0, component.y0),
            max(previous.x1, component.x1),
            max(previous.y1, component.y1),
            previous.area + component.area,
        )

    width, height = pil.size
    crops: list[tuple[int, Image.Image]] = []
    for group in groups:
        pad = max(3, min(8, round(max(group.width, group.height) * 0.12)))
        crop = pil.crop(
            (
                max(0, group.x0 - pad),
                max(0, group.y0 - pad),
                min(width, group.x1 + pad),
                min(height, group.y1 + pad),
            )
        )
        if crop.width > 2 and crop.height > 2:
            crops.append((group.x0, crop))
    return crops