noteman-slicer/noteman_slicer/detect.py

"""Detection: skew, systems, cuts, staff height.

Everything here is a *suggestion* the user confirms or edits (ADR 0004).
Nothing downstream may assume a result is right.

Systems are anchored on the vertical bracket that spans their staves, not on
gaps in the row-darkness profile: a row profile cannot tell an inter-staff gap
from an inter-system gap, and gets the count wrong on every page of a
multi-voice choral score (ADR 0006). The row profile is still needed, to expand
each anchor to its true ink extent — a bracket stops at the last staff line,
but the slice must include the lyrics printed below it.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import cv2
import numpy as np

SKEW_LIMIT_DEG = 5.0
SKEW_COARSE_STEP = 1.0
SKEW_FINE_STEP = 0.1
_SKEW_WORK_SCALE = 0.25

_INK = 128  # below this is ink, above is paper
_ANCHOR_KERNEL = 0.03  # vertical open kernel, as a fraction of page height
_ANCHOR_MIN = 0.04  # a bracket is at least this tall, as a fraction of page
_PROFILE_FLOOR = 0.02  # ink-run threshold, as a fraction of the profile peak
_EXPAND_REACH = 1.5  # how far past the bracket a system's ink reaches, in staff heights


@dataclass
class System:
    """One line of music: the ink extent that becomes a slice."""

    top: int
    bottom: int
    staff_height: float | None = None

    @property
    def height(self) -> int:
        return self.bottom - self.top


@dataclass
class PageDetection:
    skew: float
    systems: list[System] = field(default_factory=list)
    cuts: list[int] = field(default_factory=list)

    @property
    def bracketless(self) -> bool:
        """True when no bracket was found and the row profile was used alone."""
        return not self.systems or all(s.staff_height is None for s in self.systems)


def row_darkness(gray: np.ndarray) -> np.ndarray:
    return (255 - gray.astype(np.float32)).sum(axis=1)


def deskew_angle(gray: np.ndarray) -> float:
    """Angle maximising row-darkness variance — staff lines are the signal.

    Coarse then fine, on a downscaled copy: 31 warps instead of 101.
    """
    work = cv2.resize(gray, None, fx=_SKEW_WORK_SCALE, fy=_SKEW_WORK_SCALE,
                      interpolation=cv2.INTER_AREA)

    def score(angle: float) -> float:
        return float(row_darkness(_rotate(work, angle, cv2.INTER_LINEAR)).var())

    coarse = np.arange(-SKEW_LIMIT_DEG, SKEW_LIMIT_DEG + 1e-9, SKEW_COARSE_STEP)
    best = max(coarse, key=score)
    fine = np.arange(best - SKEW_COARSE_STEP, best + SKEW_COARSE_STEP + 1e-9, SKEW_FINE_STEP)
    fine = fine[np.abs(fine) <= SKEW_LIMIT_DEG]
    return round(float(max(fine, key=score)), 2)


def _rotate(gray: np.ndarray, angle: float, flags: int = cv2.INTER_CUBIC) -> np.ndarray:
    if angle == 0.0:
        return gray
    h, w = gray.shape
    m = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0)
    return cv2.warpAffine(gray, m, (w, h), flags=flags, borderValue=255)


def deskew(gray: np.ndarray, angle: float) -> np.ndarray:
    return _rotate(gray, angle)


def system_anchors(gray: np.ndarray) -> list[tuple[int, int]]:
    """y-extents of the vertical brackets, one per system."""
    h = gray.shape[0]
    binary = (gray < _INK).astype(np.uint8)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(3, int(h * _ANCHOR_KERNEL))))
    strokes = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

    count, _, stats, _ = cv2.connectedComponentsWithStats(strokes, 8)
    tall = [
        (stats[i, cv2.CC_STAT_TOP], stats[i, cv2.CC_STAT_TOP] + stats[i, cv2.CC_STAT_HEIGHT])
        for i in range(1, count)
        if stats[i, cv2.CC_STAT_HEIGHT] > h * _ANCHOR_MIN
    ]

    # Tallest first, keeping only strokes that don't overlap one already kept:
    # a system's barlines all overlap its bracket, so each system yields one.
    anchors: list[tuple[int, int]] = []
    for top, bottom in sorted(tall, key=lambda s: s[1] - s[0], reverse=True):
        if any(not (bottom < a[0] or top > a[1]) for a in anchors):
            continue
        anchors.append((top, bottom))
    return sorted(anchors)


def ink_runs(gray: np.ndarray) -> list[tuple[int, int]]:
    """Rows containing ink, despeckled — specks are the known failure mode."""
    profile = row_darkness(cv2.medianBlur(gray, 3))
    if profile.max() <= 0:
        return []
    inked = profile > profile.max() * _PROFILE_FLOOR

    runs: list[tuple[int, int]] = []
    start: int | None = None
    for i, on in enumerate(inked):
        if on and start is None:
            start = i
        elif not on and start is not None:
            runs.append((start, i))
            start = None
    if start is not None:
        runs.append((start, len(inked)))
    return runs


def staff_height(gray: np.ndarray, top: int, bottom: int) -> float | None:
    """Distance between a staff's outer lines, from staff-line spacing."""
    profile = row_darkness(gray[top:bottom])
    if profile.size == 0 or profile.max() <= 0:
        return None
    peaks = np.where(profile > profile.max() * 0.55)[0]
    if peaks.size < 2:
        return None

    centres = []
    run = [peaks[0]]
    for prev, cur in zip(peaks, peaks[1:]):
        if cur - prev > 3:
            centres.append(float(np.mean(run)))
            run = []
        run.append(cur)
    centres.append(float(np.mean(run)))
    if len(centres) < 2:
        return None

    gaps = np.diff(centres)
    # Keep intra-staff gaps; the big ones are the spaces between staves.
    intra = gaps[gaps < np.median(gaps) * 2]
    if intra.size == 0:
        return None
    return float(np.median(intra) * 4)  # 5 lines, 4 spaces


def _gap(run: tuple[int, int], span: tuple[int, int]) -> int:
    """Vertical distance between an ink run and a bracket span; 0 if they overlap."""
    start, end = run
    top, bottom = span
    if end > top and start < bottom:
        return 0
    return top - end if end <= top else start - bottom


def _assign(
    runs: list[tuple[int, int]],
    anchors: list[tuple[int, int]],
    reaches: list[float],
) -> list[tuple[int, int]]:
    """Give every ink run to one system, and return each system's extent.

    A run between two systems is resolved by **precedence, not proximity**: the
    system above wins if the run is within its reach. Text printed under a staff
    belongs to that staff, and engravers space lyrics generously — on *Feliz
    Navidad* a lyric line sits 43px under its own system's bracket but only 10px
    above the next one's, so nearest-bracket gives it to the wrong system.

    Distance is measured from the *bracket*, never from a growing extent — a
    title block's credit lines are stacked closely enough that a chaining
    expansion hops from one to the next and walks the whole way up the page.

    One pass over all systems, rather than each bracket expanding on its own, so
    that a run has exactly one owner and extents cannot overlap.

    Known limit: when a lyric line is printed tight enough under its system that
    no blank row separates it from the *next* system's staves, the two fuse into
    a single ink run and no row profile can split them — the lyric is then given
    to the system below and the cut lands high. Dragging the cut is the fix;
    separating them needs a signal this pass doesn't have.
    """
    bounds = [list(a) for a in anchors]

    def claim(index: int, run: tuple[int, int]) -> None:
        bounds[index][0] = min(bounds[index][0], run[0])
        bounds[index][1] = max(bounds[index][1], run[1])

    for run in runs:
        gaps = [_gap(run, a) for a in anchors]

        # Ink overlapping a bracket belongs to it — to the one it overlaps most,
        # whatever else is in reach.
        inside = [
            (min(run[1], anchors[i][1]) - max(run[0], anchors[i][0]), i)
            for i, g in enumerate(gaps)
            if g == 0
        ]
        if inside:
            claim(max(inside)[1], run)
            continue

        within = [i for i, g in enumerate(gaps) if g <= reaches[i]]
        if not within:
            continue  # a title block or a footer: too far from any system

        # Otherwise the system above wins, and only failing that the one below.
        above = [i for i in within if anchors[i][1] <= run[0]]
        claim(above[-1] if above else within[0], run)

    return [(lo, hi) for lo, hi in bounds]


def detect_page(gray: np.ndarray, skew: float | None = None) -> PageDetection:
    """Full proposal for one page raster. `gray` is the *unrotated* page."""
    angle = deskew_angle(gray) if skew is None else skew
    straight = deskew(gray, angle)

    runs = ink_runs(straight)
    anchors = system_anchors(straight)

    if anchors:
        # Staff height is measured on the bracket span, before expansion, so a
        # swallowed title block can't distort it.
        heights = [staff_height(straight, top, bottom) for top, bottom in anchors]
        reaches = [(h or gray.shape[0] * 0.02) * _EXPAND_REACH for h in heights]
        systems = [
            System(top=lo, bottom=hi, staff_height=h)
            for (lo, hi), h in zip(_assign(runs, anchors, reaches), heights)
        ]
    else:
        # No bracket: a single-staff melody or lead sheet, where every ink run
        # genuinely is its own system.
        systems = [System(top=t, bottom=b) for t, b in runs]

    cuts = [
        (systems[i].bottom + systems[i + 1].top) // 2 for i in range(len(systems) - 1)
    ]
    return PageDetection(skew=angle, systems=systems, cuts=cuts)