image-match/thirdparty/round_corners.py

"""给图像加圆角矩形遮罩：可导出带透明通道的 PNG，或与纯色背景合成（便于 JPG）。"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import cv2
import numpy as np

from routing import repo_path


def rounded_rect_mask(height: int, width: int, radius_px: int) -> np.ndarray:
    """单通道遮罩，圆角内为 255，外为 0。"""
    r = int(max(0, min(radius_px, min(height, width) // 2)))
    m = np.zeros((height, width), dtype=np.uint8)
    if r == 0:
        m[:] = 255
        return m

    h, w = height - 1, width - 1
    cv2.rectangle(m, (r, 0), (w - r, height - 1), 255, thickness=-1)
    cv2.rectangle(m, (0, r), (width - 1, h - r), 255, thickness=-1)

    cv2.ellipse(m, (r, r), (r, r), 180, 0, 90, 255, thickness=-1)
    cv2.ellipse(m, (w - r, r), (r, r), 270, 0, 90, 255, thickness=-1)
    cv2.ellipse(m, (w - r, h - r), (r, r), 0, 0, 90, 255, thickness=-1)
    cv2.ellipse(m, (r, h - r), (r, r), 90, 0, 90, 255, thickness=-1)
    return m


def round_corners(
    input_path: str | Path,
    output_path: str | Path,
    radius_px: int,
    *,
    background_bgr: tuple[int, int, int] | None = None,
    feather_px: float = 0.0,
) -> Path:
    """
    :param background_bgr: 若为 ``None``，输出保留透明（圆角外 alpha=0）；否则与给定 BGR 底色合成（适合 ``.jpg``）。
    :param feather_px: 对遮罩做轻微高斯模糊半径（像素），边缘更柔和；``0`` 为硬边。
    """
    inp = repo_path(input_path)
    out = repo_path(output_path)

    img = cv2.imread(str(inp), cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f"无法读取：{inp}")

    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    h, w = img.shape[:2]
    mask = rounded_rect_mask(h, w, radius_px).astype(np.float32) / 255.0

    if feather_px > 0:
        k = max(1, int(round(feather_px * 2)) | 1)
        mask_u8 = (mask * 255).astype(np.uint8)
        mask_u8 = cv2.GaussianBlur(mask_u8, (k, k), feather_px)
        mask = mask_u8.astype(np.float32) / 255.0

    out.parent.mkdir(parents=True, exist_ok=True)

    if background_bgr is None:
        if img.shape[2] == 4:
            bgra = img.astype(np.float32)
            a_in = bgra[:, :, 3:4] / 255.0
            a_out = mask[..., np.newaxis] * a_in
            bgra[:, :, :3] *= mask[..., np.newaxis]
            bgra[:, :, 3:4] = np.clip(a_out * 255.0, 0, 255)
            write_img = np.clip(bgra, 0, 255).astype(np.uint8)
        else:
            bgr = img.astype(np.float32)
            bgr *= mask[..., np.newaxis]
            alpha = (mask * 255.0).astype(np.uint8)
            write_img = cv2.merge(
                [
                    bgr[:, :, 0].astype(np.uint8),
                    bgr[:, :, 1].astype(np.uint8),
                    bgr[:, :, 2].astype(np.uint8),
                    alpha,
                ]
            )
        suffix = out.suffix.lower()
        if suffix not in (".png", ".webp", ".tif", ".tiff"):
            out = out.with_suffix(".png")
        if not cv2.imwrite(str(out), write_img):
            raise OSError(f"无法写入：{out}")
        return out

    bg_b, bg_g, bg_r = background_bgr
    if img.shape[2] == 4:
        bgr = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR).astype(np.float32)
        alpha = (img[:, :, 3].astype(np.float32) / 255.0) * mask
    else:
        bgr = img.astype(np.float32)
        alpha = mask

    bg = np.zeros_like(bgr)
    bg[:, :, 0] = bg_b
    bg[:, :, 1] = bg_g
    bg[:, :, 2] = bg_r
    a = alpha[..., np.newaxis]
    blended = np.clip(bgr * a + bg * (1.0 - a), 0, 255).astype(np.uint8)
    if not cv2.imwrite(str(out), blended):
        raise OSError(f"无法写入：{out}")
    return out


def round_corner(
    *,
    degree: int,
    image: str | Path,
    output_path: str | Path,
    background_bgr: tuple[int, int, int] | None = None,
    feather_px: float = 0.0,
) -> Path:
    """
    与 ``main`` 等调用约定一致的关键字接口。

    ``degree``：圆角半径（像素），取值含义与 ``round_corners(..., radius_px=...)`` 相同（命名为 ``degree`` 仅为兼容现有调用）。
    ``image`` / ``output_path``：相对路径相对仓库根目录（见 ``routing.repo_path``）。
    """
    return round_corners(
        image,
        output_path,
        int(degree),
        background_bgr=background_bgr,
        feather_px=feather_px,
    )


def main() -> None:
    ap = argparse.ArgumentParser(description="给图片加圆角并保存（相对路径相对仓库根目录）。")
    ap.add_argument("input", type=Path, help="输入图路径")
    ap.add_argument("-o", "--output", type=Path, required=True, help="输出路径")
    ap.add_argument(
        "-r",
        "--radius",
        type=int,
        default=24,
        help="圆角半径（像素），不超过短边一半",
    )
    ap.add_argument(
        "--bg",
        type=str,
        default=None,
        help="合成底色 B,G,R（0-255），逗号分隔；省略则输出透明 PNG（自动改后缀为 .png）",
    )
    ap.add_argument(
        "--feather",
        type=float,
        default=0.0,
        help="遮罩边缘羽化（高斯 sigma，像素），0 为硬边",
    )
    args = ap.parse_args()

    bg = None
    if args.bg:
        parts = [int(x.strip()) for x in args.bg.split(",")]
        if len(parts) != 3:
            raise SystemExit("--bg 须为三个整数：B,G,R")
        bg = (parts[0], parts[1], parts[2])

    p = round_corners(
        args.input,
        args.output,
        args.radius,
        background_bgr=bg,
        feather_px=args.feather,
    )
    print(f"已写入：{p}", flush=True)


if __name__ == "__main__":
    if sys.platform == "win32":
        for stream in (sys.stdout, sys.stderr):
            try:
                stream.reconfigure(encoding="utf-8")
            except Exception:
                pass
    main()