app/src/inkference/htr/segmentation.py

"""Line segmentation — Kraken baseline segmentation (`blla`).

Faithful port of notebooks/segmentManualTranscriptionsIntoLinesWithKraken*.ipynb:
Kraken's bundled baseline model predicts, per line, a **baseline + boundary
polygon**; we take the polygon, derive a bbox, drop noise boxes (MIN_W/MIN_H),
sort top-to-bottom by vertical centre, and crop each line **masked to its polygon**
(dilated by POLY_PAD) so neighbouring ascenders/descenders don't bleed in.

The old dependency-light ProjectionSegmenter is disabled — see the commented block
at the bottom. Kraken is now the only backend.

Heavy imports (kraken/torch/PIL) are lazy so this module imports with nothing installed.
"""
from __future__ import annotations

import importlib.util
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Protocol

from ..config import HTRConfig
from ..config import htr as default_htr
from ..schemas import BBox

if TYPE_CHECKING:  # pragma: no cover
    from PIL import Image


@dataclass
class LineCrop:
    """One detected line: its polygon-masked crop plus where it sits on the page."""

    index: int
    bbox: BBox
    image: "Image.Image"
    polygon: list[tuple[float, float]] = field(default_factory=list)


class Segmenter(Protocol):
    def segment(self, page: "Image.Image") -> list[LineCrop]: ...


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def downscale(page: "Image.Image", max_long_edge: int) -> tuple["Image.Image", float]:
    """Cap the long edge to keep free-CPU work bounded. Returns (image, scale)."""
    w, h = page.size
    long_edge = max(w, h)
    if long_edge <= max_long_edge:
        return page, 1.0
    scale = max_long_edge / long_edge
    return page.resize((round(w * scale), round(h * scale))), scale


def _resolve_device(device: str) -> str:
    if device != "auto":
        return device
    try:
        import torch

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:  # pragma: no cover
        return "cpu"


def _boundary_of(line):
    """A line's boundary polygon as a list of points, across kraken API versions."""
    b = getattr(line, "boundary", None)
    if b is None and isinstance(line, dict):
        b = line.get("boundary")
    return b


# --------------------------------------------------------------------------- #
# Kraken backend
# --------------------------------------------------------------------------- #
class KrakenSegmenter:
    """Kraken baseline segmentation with polygon-masked line crops."""

    def __init__(self, cfg: HTRConfig = default_htr) -> None:
        self.cfg = cfg
        self.device = _resolve_device(cfg.device)
        self._seg_model = None
        self._model_loaded = False

    def _ensure_model(self):
        """Load Kraken's bundled default baseline model (blla.mlmodel) once.

        Falls back to None, in which case blla.segment uses its own default."""
        if self._model_loaded:
            return self._seg_model
        import os

        import kraken
        from kraken.lib import vgsl

        cand = os.path.join(os.path.dirname(kraken.__file__), "blla.mlmodel")
        if os.path.exists(cand):
            self._seg_model = vgsl.TorchVGSLModel.load_model(cand)
        else:
            self._seg_model = None
        self._model_loaded = True
        return self._seg_model

    def _run_blla(self, im):
        from kraken import blla

        model = self._ensure_model()
        # the `device` kwarg exists on newer kraken; fall back gracefully.
        try:
            return blla.segment(im, model=model, device=self.device)
        except TypeError:
            return blla.segment(im, model=model) if model is not None else blla.segment(im)

    def segment(self, page: "Image.Image") -> list[LineCrop]:
        from PIL import Image, ImageDraw, ImageFilter

        rgb = page.convert("RGB")
        seg = self._run_blla(rgb)
        lines = getattr(seg, "lines", None)
        if lines is None and isinstance(seg, dict):
            lines = seg.get("lines", [])
        lines = lines or []

        W, H = rgb.size
        cfg = self.cfg

        # 1. polygon -> bbox record, dropping noise boxes (MIN_W / MIN_H).
        records: list[tuple[BBox, list[tuple[float, float]]]] = []
        for ln in lines:
            poly = _boundary_of(ln)
            if not poly:
                continue
            pts = [(float(p[0]), float(p[1])) for p in poly]
            xs = [p[0] for p in pts]
            ys = [p[1] for p in pts]
            bbox = (min(xs), min(ys), max(xs), max(ys))
            if (bbox[2] - bbox[0]) < cfg.min_w or (bbox[3] - bbox[1]) < cfg.min_h:
                continue
            records.append((bbox, pts))

        # 2. reading order: top-to-bottom by vertical centre.
        records.sort(key=lambda r: (r[0][1] + r[0][3]) / 2)

        # 3. polygon-masked crop per line (PAD_X/PAD_Y + MaxFilter(POLY_PAD) dilation).
        crops: list[LineCrop] = []
        for idx, (bbox, pts) in enumerate(records):
            x0, y0, x1, y1 = bbox
            cx0 = max(0, int(x0) - cfg.pad_x)
            cy0 = max(0, int(y0) - cfg.pad_y)
            cx1 = min(W, int(x1) + cfg.pad_x)
            cy1 = min(H, int(y1) + cfg.pad_y)
            crop = rgb.crop((cx0, cy0, cx1, cy1))

            if cfg.mask_to_polygon:
                mask = Image.new("L", crop.size, 0)
                ImageDraw.Draw(mask).polygon(
                    [(px - cx0, py - cy0) for px, py in pts], fill=255
                )
                if cfg.poly_pad > 0:  # dilate so the line's own strokes aren't clipped
                    mask = mask.filter(ImageFilter.MaxFilter(cfg.poly_pad * 2 + 1))
                white = Image.new("RGB", crop.size, (255, 255, 255))
                crop = Image.composite(crop, white, mask)

            crops.append(
                LineCrop(index=idx, bbox=(cx0, cy0, cx1, cy1), image=crop, polygon=pts)
            )
        return crops


def kraken_available() -> bool:
    return importlib.util.find_spec("kraken") is not None


def get_segmenter(name: str = "kraken", cfg: HTRConfig = default_htr) -> Segmenter:
    """Return the Kraken segmenter. The projection fallback is disabled."""
    if name in ("kraken", "auto"):
        if not kraken_available():
            raise ImportError(
                "Kraken is not installed. Install it (see the note in requirements.txt — "
                "Kraken pins its own torch, so a dedicated venv is recommended)."
            )
        return KrakenSegmenter(cfg)
    if name == "projection":
        raise ValueError("ProjectionSegmenter is disabled; use the Kraken segmenter.")
    raise ValueError(f"unknown segmenter: {name!r}")


# --------------------------------------------------------------------------- #
# DISABLED — dependency-light projection-profile fallback (kept for reference).
# Re-enable in get_segmenter() only if you deliberately want a Kraken-free path.
# --------------------------------------------------------------------------- #
# class ProjectionSegmenter:
#     """Find line bands from the horizontal ink-projection profile. Weak on dense
#     or multi-column layouts — replaced by Kraken baseline segmentation."""
#
#     def __init__(self, cfg: HTRConfig = default_htr) -> None:
#         self.cfg = cfg
#
#     def segment(self, page: "Image.Image") -> list[LineCrop]:
#         import numpy as np
#         from PIL import ImageOps
#
#         gray = ImageOps.grayscale(page)
#         arr = np.asarray(gray, dtype=np.float32)
#         h, w = arr.shape
#         thresh = float(arr.mean()) - 0.4 * float(arr.std())
#         ink = (arr < thresh).astype(np.float32)
#         row_ink = ink.sum(axis=1)
#         if row_ink.max() <= 0:
#             return []
#         active = row_ink > (0.04 * row_ink.max())
#         bands: list[tuple[int, int]] = []
#         start = None
#         for y in range(h):
#             if active[y] and start is None:
#                 start = y
#             elif not active[y] and start is not None:
#                 bands.append((start, y))
#                 start = None
#         if start is not None:
#             bands.append((start, h))
#         rgb = page.convert("RGB")
#         crops: list[LineCrop] = []
#         idx = 0
#         for (y0, y1) in bands:
#             if (y1 - y0) < self.cfg.min_h:
#                 continue
#             band_ink = ink[y0:y1].sum(axis=0)
#             cols = np.where(band_ink > 0)[0]
#             if cols.size == 0:
#                 continue
#             x0, x1 = int(cols[0]), int(cols[-1]) + 1
#             if (x1 - x0) < self.cfg.min_w:
#                 continue
#             x0p = max(0, x0 - self.cfg.pad_x); y0p = max(0, y0 - self.cfg.pad_y)
#             x1p = min(w, x1 + self.cfg.pad_x); y1p = min(h, y1 + self.cfg.pad_y)
#             crops.append(LineCrop(index=idx, bbox=(x0p, y0p, x1p, y1p),
#                                   image=rgb.crop((x0p, y0p, x1p, y1p))))
#             idx += 1
#         return crops