image-match/thirdparty/crop_img.py

"""
按四边比例裁剪图像（路径或 ``numpy`` ndarray）。模块：``thirdparty/crop_img.py``。

文件读入默认 **保留颜色**（``IMREAD_UNCHANGED``：灰度 / BGR / BGRA），**不强行转灰度**。

**只做裁剪，不做缩放**：内部仅为 ``numpy`` 切片 ``img[y0:y1, x0:x1]``，**从不调用** ``cv2.resize`` 等插值。
若原图接近正方形、且只裁上下不裁左右，则输出 **像素宽度与原版相同**、**高度变矮**，宽高比变大，
在显示器上会显得「横向更宽」——这是比例变化，不是把图横向拉伸。

四个比例参数：取值 **0～1**，直接表示该边要裁掉的相对长度——
上、下相对**图像高度**，左、右相对**图像宽度**；**1.0 = 裁掉该边方向 100%**。
例如 ``0.1, 0.1, 0, 0`` 表示上下各裁 10% 高；``0.1`` 即 10%。总裁切过量时会压缩，保证至少剩 4×4 像素。

可选第 6 个位置参数 ``output_path``：若给出（相对路径相对仓库根目录），则保存到该路径并等价于 ``save=True``。

否则由关键字 ``save``、``output_dir``、``output_filename`` 控制写出（默认 ``output/cut-img.png``）。

简约用法（顺序：上、下、左、右、图源、可选写出路径）::

    CropImg().crop_edges(0.1, 0.1, 0.0, 0.0, path, "output/crop-img.png")
    CropImg().crop_edges(0.1, 0.1, 0.0, 0.0, path, save=True, output_dir=..., output_filename="...")
"""

from __future__ import annotations

from pathlib import Path
from typing import Union

import cv2
import numpy as np

from routing import REPO_ROOT, repo_path

ImageSource = Union[str, Path, np.ndarray]

_DEFAULT_OUTPUT_NAME = "cut-img.png"


def read_image_unchanged(path: str | Path) -> np.ndarray:
    """
    与裁剪管线读图一致：``IMREAD_UNCHANGED``，二值/灰度为 ``(H,W)``，彩色为 ``(H,W,3)`` 或带透明 ``(H,W,4)``。
    """
    p = repo_path(path)
    img = cv2.imread(str(p), cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f"无法读取图片：{p}")
    return img


def _load_for_crop(image: ImageSource) -> np.ndarray:
    if isinstance(image, np.ndarray):
        arr = image
        if arr.ndim == 2:
            return arr
        if arr.ndim == 3 and arr.shape[2] in (3, 4):
            return arr
        raise ValueError("ndarray 须为 (H,W) 或 (H,W,3)/(H,W,4)。")
    return read_image_unchanged(image)


def _crop_frac(x: float | int) -> float:
    """裁切比例：0～1，1 为 100% 边长。"""
    return float(np.clip(float(x), 0.0, 1.0))


def crop_slice_bounds(
    h: int,
    w: int,
    top: float | int,
    bottom: float | int,
    left: float | int,
    right: float | int,
) -> tuple[int, int, int, int]:
    """
    与 ``crop_image_edges`` 相同规则，返回 ``img[y0:y1, x0:x1]`` 的半开切片（与 numpy 一致）。

    返回值 ``(y0, y1, x0, x1)``；若无法安全裁剪则 ``(0, h, 0, w)`` 表示整幅。
    """
    if h < 4 or w < 4:
        return 0, h, 0, w

    tf = _crop_frac(top)
    bf = _crop_frac(bottom)
    lf = _crop_frac(left)
    rf = _crop_frac(right)

    dt = int(round(h * tf))
    db = int(round(h * bf))
    dl = int(round(w * lf))
    dr = int(round(w * rf))
    dt, db = max(0, dt), max(0, db)
    dl, dr = max(0, dl), max(0, dr)
    if dt + db > h - 4:
        s = (h - 4) / max(dt + db, 1)
        dt = int(dt * s)
        db = int(db * s)
        db = (h - 4) - dt
    if dl + dr > w - 4:
        s = (w - 4) / max(dl + dr, 1)
        dl = int(dl * s)
        dr = int(dr * s)
        dr = (w - 4) - dl

    nh = h - dt - db
    nw = w - dl - dr
    if nh < 4 or nw < 4:
        return 0, h, 0, w
    y0, y1, x0, x1 = dt, h - db, dl, w - dr
    return y0, y1, x0, x1


def crop_image_edges(
    top: float | int,
    bottom: float | int,
    left: float | int,
    right: float | int,
    image: ImageSource,
    output_path: str | Path | None = None,
    *,
    save: bool = True,
    output_dir: str | Path | None = None,
    output_filename: str = _DEFAULT_OUTPUT_NAME,
) -> np.ndarray:
    """
    按上、下、左、右四边裁切；各参数为 **0～1**，即该边裁掉的相对高/宽（``0.1`` = 10%）。

    :param top, bottom: 相对**高度**的比例
    :param left, right: 相对**宽度**的比例
    :param image: 路径，或 ``(H,W)`` / ``(H,W,3)`` / ``(H,W,4)`` ndarray（相对路径相对仓库根目录）
    :param output_path: 若给定则写入该文件并强制 ``save=True``（目录由路径决定）
    :param save: 未给 ``output_path`` 时是否写入 ``output_dir / output_filename``
    """
    if output_path is not None:
        out_full = repo_path(output_path)
        output_dir = out_full.parent
        output_filename = out_full.name
        save = True

    img = _load_for_crop(image)
    h, w = int(img.shape[0]), int(img.shape[1])
    y0, y1, x0, x1 = crop_slice_bounds(h, w, top, bottom, left, right)
    if (y0, y1, x0, x1) == (0, h, 0, w):
        out = img
    else:
        # 纯切片裁切；不改分辨率、不做插值（与 cv2.resize 无关）。
        out = np.ascontiguousarray(img[y0:y1, x0:x1])

    if save:
        _save_cropped_image(out, output_dir, output_filename)
    return out


def _save_cropped_image(
    image: np.ndarray,
    output_dir: str | Path | None,
    output_filename: str,
) -> Path:
    if output_dir is not None:
        out_dir = repo_path(output_dir)
    else:
        out_dir = REPO_ROOT / "output"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / output_filename
    if not cv2.imwrite(str(out_path), image):
        raise OSError(f"无法写入：{out_path.resolve()}")
    return out_path


def write_gray(path: str | Path, gray: np.ndarray) -> Path:
    """写入灰度图（路径后缀决定格式，``.png`` 最稳妥）；相对路径相对仓库根目录。"""
    out = repo_path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    if not cv2.imwrite(str(out), gray):
        raise OSError(f"无法写入：{out}")
    return out


class CropImg:
    """四边比例裁剪；无状态，可重复 ``CropImg().crop_edges(...)``。"""

    def crop_edges(
        self,
        top: float | int,
        bottom: float | int,
        left: float | int,
        right: float | int,
        image: ImageSource,
        output_path: str | Path | None = None,
        *,
        save: bool = True,
        output_dir: str | Path | None = None,
        output_filename: str = _DEFAULT_OUTPUT_NAME,
    ) -> np.ndarray:
        return crop_image_edges(
            top,
            bottom,
            left,
            right,
            image,
            output_path=output_path,
            save=save,
            output_dir=output_dir,
            output_filename=output_filename,
        )