staff_rectify.py

"""
Single-staff dewarp: sample mid-staff ink curve along x, fit a low-degree polynomial,
apply vertical remap so the curve becomes flat (OpenCV remap).
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np


def _column_mid_peak(
    horizontal: np.ndarray,
    x: int,
    line_ys_local: List[int],
) -> Optional[int]:
    """One y per column near expected staff band (middle line anchor)."""
    h = horizontal.shape[0]
    if h < 8:
        return None
    mid = int(round((line_ys_local[0] + line_ys_local[-1]) / 2.0))
    span = max(int(line_ys_local[-1] - line_ys_local[0]) * 2, h // 3)
    y0 = max(0, mid - span)
    y1 = min(h, mid + span)
    col = horizontal[y0:y1, x]
    if col.size == 0 or float(np.max(col)) < 1.0:
        return None
    return int(y0 + int(np.argmax(col)))


def rectify_staff_crop_bgr(
    work_bgr: np.ndarray,
    staff_x0: int,
    staff_y0: int,
    staff_w: int,
    staff_h: int,
    line_ys_global: List[int],
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Crop staff from work image, estimate vertical curvature of staff ink, remap to flatten.

    Returns (rectified_bgr_crop, meta) where meta includes staff_dewarp status.
    line_ys_global: five staff line y positions in full work image coordinates.
    """
    meta: Dict[str, Any] = {
        "staff_dewarp": "skipped",
        "staff_dewarp_model": None,
        "staff_dewarp_detail": None,
    }

    h_img, w_img = work_bgr.shape[:2]
    x0 = max(0, int(staff_x0))
    y0 = max(0, int(staff_y0))
    x1 = min(w_img, x0 + max(1, int(staff_w)))
    y1 = min(h_img, y0 + max(1, int(staff_h)))
    roi = work_bgr[y0:y1, x0:x1].copy()
    if roi.size == 0:
        return roi, meta

    line_local = [int(round(y - y0)) for y in line_ys_global]
    if len(line_local) < 5 or min(line_local) < 0 or max(line_local) >= roi.shape[0]:
        meta["staff_dewarp_detail"] = "invalid_line_ys"
        return roi, meta

    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    binary = cv2.adaptiveThreshold(
        blur,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        25,
        11,
    )
    kw = max(15, roi.shape[1] // 24)
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kw, 1))
    horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=1)

    w = roi.shape[1]
    step = max(1, w // 120)
    xs_list: List[float] = []
    ys_list: List[float] = []
    for x in range(0, w, step):
        yp = _column_mid_peak(horizontal, x, line_local)
        if yp is not None:
            xs_list.append(float(x))
            ys_list.append(float(yp))

    if len(xs_list) < 6:
        meta["staff_dewarp_detail"] = "too_few_samples"
        return roi, meta

    xs = np.array(xs_list, dtype=np.float64)
    ys = np.array(ys_list, dtype=np.float64)
    deg = 2 if len(xs_list) >= 8 else 1
    try:
        coeffs = np.polyfit(xs, ys, deg)
    except (np.linalg.LinAlgError, ValueError):
        meta["staff_dewarp_detail"] = "polyfit_failed"
        return roi, meta

    h, w = roi.shape[:2]
    x_grid = np.arange(w, dtype=np.float64)
    curve = np.polyval(coeffs, x_grid)
    c = float(np.mean(curve))
    delta = (curve - c).astype(np.float32)
    map_x = np.tile(np.arange(w, dtype=np.float32), (h, 1))
    map_y = np.arange(h, dtype=np.float32)[:, None] + delta[None, :]

    out = cv2.remap(roi, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    meta["staff_dewarp"] = "applied"
    meta["staff_dewarp_model"] = f"poly{deg}"
    meta["staff_dewarp_detail"] = "midline_peak_column_samples"
    return out, meta