xhs-note-crawling/workplace/ocr-pos.py

#!/usr/bin/env python3
"""
使用仓库内 ``python/RapidOCR`` 对**图片**做 OCR。

``ocr_find_text_center``：在单行识别结果中找完整包含目标子串的框（与 CLI 一致）。

``ocr_find_text_center_allowing_note_title_wrapped_across_two_ocr_lines``（笔记标题锚点）：
仅使用用户字符串 **第一逻辑行**（首个换行符之前）；第二行及以后 **不参与** 匹配与几何计算。
对第一行先做 **NFKC**，再抠掉难以 OCR 的字符得到多段可识别文本。整图 **只跑一次** OCR，
每段取面积最小匹配框中心；同一段内多段则取首末段中心的中点作为代表点（不再跨多逻辑行拼接）。

与 Playwright / 浏览器 DOM 无关；仅输入图片路径或字节。

用法：``python workplace/ocr-pos.py <目标文字> <图片路径>``
成功时 stdout 一行：``cx cy``（整数）。

可调模块级常量：``OCR_RAPIDOCR_GLOBAL_MAX_SIDE_LEN``、``OCR_RAPIDOCR_INFERENCE_ENGINE_NAME``、
``OCR_RAPIDOCR_ONNX_RUNTIME_USE_CUDA``（仅 ``onnxruntime`` 引擎时对应 ``EngineConfig.onnxruntime.use_cuda``）。
可选推理依赖见 ``python/RapidOCR/requirements-inference-*.txt``。
"""

from __future__ import annotations

import sys
import unicodedata
from pathlib import Path
from typing import Any

import cv2
import numpy as np

_REPO_ROOT = Path(__file__).resolve().parent.parent
_RAPIDOCR_PKG_ROOT = _REPO_ROOT / "python" / "RapidOCR" / "python"
if str(_RAPIDOCR_PKG_ROOT) not in sys.path:
    sys.path.insert(0, str(_RAPIDOCR_PKG_ROOT))

REPO_ROOT: Path = _REPO_ROOT

# OCR_RAPIDOCR_GLOBAL_MAX_SIDE_LEN: int = 2000  # 提速可试 1280 或 960（顶栏 OCR 常够用；过小易丢小字）
# OCR_RAPIDOCR_INFERENCE_ENGINE_NAME: str = "onnxruntime"  # 可选 openvino（Intel CPU 常较快）/ tensorrt（需 NVIDIA 环境）
# OCR_RAPIDOCR_ONNX_RUNTIME_USE_CUDA: bool = False  # 已装 onnxruntime-gpu 且要用 GPU 时改为 True

OCR_RAPIDOCR_GLOBAL_MAX_SIDE_LEN: int = 2000  # 提速可试 1280 或 960（顶栏 OCR 常够用；过小易丢小字）
OCR_RAPIDOCR_INFERENCE_ENGINE_NAME: str = "onnxruntime"  # 与 requirements 中 onnxruntime 一致；本机已装 openvino 可改 openvino
OCR_RAPIDOCR_ONNX_RUNTIME_USE_CUDA: bool = False  # 已装 onnxruntime-gpu 且要用 GPU 时改为 True


# 下方为 RapidOCR 单例缓存（勿改此变量）。可切换的只有上面三个常量；改完后须重启进程才会重新建引擎。
_ocr_engine_singleton = None


def _build_rapid_ocr_params_dictionary() -> dict[str, Any]:
    from rapidocr.utils.typings import EngineType

    inference_engine_name_normalized = (
        OCR_RAPIDOCR_INFERENCE_ENGINE_NAME.strip().lower()
    )
    rapid_ocr_params: dict[str, Any] = {
        "Global.max_side_len": int(OCR_RAPIDOCR_GLOBAL_MAX_SIDE_LEN),
        "Global.use_cls": False,
    }

    if inference_engine_name_normalized == "openvino":
        rapid_ocr_params["Det.engine_type"] = EngineType.OPENVINO
        rapid_ocr_params["Cls.engine_type"] = EngineType.OPENVINO
        rapid_ocr_params["Rec.engine_type"] = EngineType.OPENVINO
        return rapid_ocr_params

    if inference_engine_name_normalized == "tensorrt":
        rapid_ocr_params["Det.engine_type"] = EngineType.TENSORRT
        rapid_ocr_params["Cls.engine_type"] = EngineType.TENSORRT
        rapid_ocr_params["Rec.engine_type"] = EngineType.TENSORRT
        return rapid_ocr_params

    rapid_ocr_params["Det.engine_type"] = EngineType.ONNXRUNTIME
    rapid_ocr_params["Cls.engine_type"] = EngineType.ONNXRUNTIME
    rapid_ocr_params["Rec.engine_type"] = EngineType.ONNXRUNTIME
    rapid_ocr_params["EngineConfig.onnxruntime.use_cuda"] = bool(
        OCR_RAPIDOCR_ONNX_RUNTIME_USE_CUDA
    )
    return rapid_ocr_params


def _get_rapid_ocr():
    global _ocr_engine_singleton
    if _ocr_engine_singleton is None:
        from rapidocr import RapidOCR

        _ocr_engine_singleton = RapidOCR(
            params=_build_rapid_ocr_params_dictionary(),
        )
    return _ocr_engine_singleton


def _rapidocr_input_from_path_bytes_or_bgr_numpy(
    image: str | Path | bytes | np.ndarray,
):
    if isinstance(image, np.ndarray):
        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


def _quad_center_xy(box) -> tuple[float, float]:
    """
    ``box`` 为四角点 (4,2)；返回几何中心点。
    """
    xs = [float(p[0]) for p in box]
    ys = [float(p[1]) for p in box]
    return sum(xs) / len(xs), sum(ys) / len(ys)


def _normalize_ocr_match_string(s: str) -> str:
    """去掉首尾空白并去掉内部空白；再将全角标点与半角对齐，便于 metadata 与 RapidOCR 输出对齐。"""
    collapsed_whitespace_string = "".join((s or "").strip().split())
    punctuation_aligned_string = (
        collapsed_whitespace_string.replace("，", ",")
        .replace("？", "?")
        .replace("（", "(")
        .replace("）", ")")
        .replace("：", ":")
        .replace("；", ";")
        .replace("！", "!")
    )
    return punctuation_aligned_string.casefold()


def _ocr_line_contains_full_target(ocr_line: str, target_text: str) -> bool:
    """
    仅当 **整段** ``target_text`` 连续出现在本行识别结果中才算匹配。

    禁止再用「识别串是目标子串」（如仅「小红书」）去碰「搜索小红书」，
    避免中心点落到不完整的字块上（例如「小红书上」里截出的「小红书」）。
    """
    line = _normalize_ocr_match_string(ocr_line)
    target = _normalize_ocr_match_string(target_text)
    if not target:
        return False
    return target in line


def _ocr_dom_anchor_line_matches_normalized_text(
    ocr_line: str,
    anchor_norm: str,
) -> bool:
    """与 ``clear-input`` 一致：多字用子串匹配，单字要求整行一致。"""
    norm = _normalize_ocr_match_string(ocr_line)
    if not anchor_norm:
        return False
    if len(anchor_norm) >= 2:
        return anchor_norm in norm
    return norm == anchor_norm


def ocr_smallest_line_center_for_dom_text(
    image: str | Path | bytes | np.ndarray,
    anchor_dom_text: str,
) -> tuple[int, int] | None:
    """
    在 OCR 结果中找归一后包含 ``anchor_dom_text`` 的行；多行命中时取**框面积最小**的一行，
    返回该行中心（屏幕/图像像素）。与 ``clear-input`` 锚点规则一致。
    """
    anchor_norm = _normalize_ocr_match_string(anchor_dom_text)
    if not anchor_norm:
        return None
    ocr = _get_rapid_ocr()
    result = ocr(_rapidocr_input_from_path_bytes_or_bgr_numpy(image))
    if result is None or not result.txts or result.boxes is None:
        return None
    anchor_indices: list[int] = []
    for i, txt in enumerate(result.txts):
        if _ocr_dom_anchor_line_matches_normalized_text(txt, anchor_norm):
            anchor_indices.append(i)
    if not anchor_indices:
        return None

    def _box_area(idx: int) -> float:
        box = result.boxes[idx]
        xs = [float(p[0]) for p in box]
        ys = [float(p[1]) for p in box]
        return (max(xs) - min(xs)) * (max(ys) - min(ys))

    best_i = min(anchor_indices, key=_box_area)
    cx, cy = _quad_center_xy(result.boxes[best_i])
    return int(round(cx)), int(round(cy))


_CJK_AND_COMMON_TITLE_PUNCTUATION_CHARACTER_FROZEN_SET = frozenset(
    "，。、：；？！""''（）《》·—…－—「」『』【】〈〉＆％＃＠…—",
)


def _is_single_character_treated_as_recognizable_by_rapidocr_line_text_anchor(
    single_unicode_character: str,
) -> bool:
    if len(single_unicode_character) != 1:
        return False
    code_point_integer = ord(single_unicode_character)
    if single_unicode_character.isascii():
        return single_unicode_character.isalnum() or single_unicode_character in " \t.-_/:&!?"
    if (
        0x4E00 <= code_point_integer <= 0x9FFF
        or 0x3400 <= code_point_integer <= 0x4DBF
        or 0x20000 <= code_point_integer <= 0x2CEAF
    ):
        return True
    if single_unicode_character in _CJK_AND_COMMON_TITLE_PUNCTUATION_CHARACTER_FROZEN_SET:
        return True
    return False


def _split_user_anchor_string_into_recognizable_text_segment_string_list(
    raw_anchor_string_from_user: str,
) -> list[str]:
    unicode_nfkc_normalized_anchor_string = unicodedata.normalize(
        "NFKC",
        raw_anchor_string_from_user or "",
    )
    contiguous_recognizable_character_list: list[str] = []
    recognizable_text_segment_string_list: list[str] = []
    for single_character in unicode_nfkc_normalized_anchor_string:
        if _is_single_character_treated_as_recognizable_by_rapidocr_line_text_anchor(
            single_character,
        ):
            contiguous_recognizable_character_list.append(single_character)
        else:
            if contiguous_recognizable_character_list:
                recognizable_text_segment_string_list.append(
                    "".join(contiguous_recognizable_character_list),
                )
                contiguous_recognizable_character_list = []
    if contiguous_recognizable_character_list:
        recognizable_text_segment_string_list.append(
            "".join(contiguous_recognizable_character_list),
        )
    return recognizable_text_segment_string_list


def _split_anchor_raw_string_into_logical_line_then_recognizable_segment_nested_string_list(
    raw_anchor_string_from_user: str,
) -> list[list[str]]:
    anchor_string_with_normalized_newline_only = (
        (raw_anchor_string_from_user or "")
        .replace("\r\n", "\n")
        .replace("\r", "\n")
    )
    logical_line_raw_string_list = anchor_string_with_normalized_newline_only.split("\n")
    nested_recognizable_segment_string_list: list[list[str]] = []
    for single_logical_line_raw_string in logical_line_raw_string_list:
        stripped_logical_line_string = single_logical_line_raw_string.strip()
        if not stripped_logical_line_string:
            continue
        recognizable_segment_list_for_one_logical_line = (
            _split_user_anchor_string_into_recognizable_text_segment_string_list(
                stripped_logical_line_string,
            )
        )
        if recognizable_segment_list_for_one_logical_line:
            nested_recognizable_segment_string_list.append(
                recognizable_segment_list_for_one_logical_line,
            )
    return nested_recognizable_segment_string_list


def _midpoint_float_xy_pair_from_first_and_last_center_in_ordered_float_pair_list(
    ordered_center_xy_float_pair_list: list[tuple[float, float]],
) -> tuple[float, float] | None:
    if not ordered_center_xy_float_pair_list:
        return None
    if len(ordered_center_xy_float_pair_list) == 1:
        lone_center_x, lone_center_y = ordered_center_xy_float_pair_list[0]
        return lone_center_x, lone_center_y
    first_center_x, first_center_y = ordered_center_xy_float_pair_list[0]
    last_center_x, last_center_y = ordered_center_xy_float_pair_list[-1]
    return (
        (first_center_x + last_center_x) / 2.0,
        (first_center_y + last_center_y) / 2.0,
    )


def _midpoint_float_xy_pair_from_closest_pair_of_points_in_float_pair_list(
    center_xy_float_pair_list: list[tuple[float, float]],
) -> tuple[float, float] | None:
    point_count = len(center_xy_float_pair_list)
    if point_count == 0:
        return None
    if point_count == 1:
        lone_x, lone_y = center_xy_float_pair_list[0]
        return lone_x, lone_y
    best_squared_euclidean_distance_between_pair: float | None = None
    midpoint_x_between_closest_pair = 0.0
    midpoint_y_between_closest_pair = 0.0
    for first_point_index in range(point_count):
        first_x, first_y = center_xy_float_pair_list[first_point_index]
        for second_point_index in range(first_point_index + 1, point_count):
            second_x, second_y = center_xy_float_pair_list[second_point_index]
            delta_x = first_x - second_x
            delta_y = first_y - second_y
            squared_distance = delta_x * delta_x + delta_y * delta_y
            if (
                best_squared_euclidean_distance_between_pair is None
                or squared_distance < best_squared_euclidean_distance_between_pair
            ):
                best_squared_euclidean_distance_between_pair = squared_distance
                midpoint_x_between_closest_pair = (first_x + second_x) / 2.0
                midpoint_y_between_closest_pair = (first_y + second_y) / 2.0
    return midpoint_x_between_closest_pair, midpoint_y_between_closest_pair


def _ocr_smallest_area_line_index_where_normalized_line_contains_substring(
    ocr_result: Any,
    substring_normalized: str,
) -> int | None:
    if not substring_normalized:
        return None
    best_line_index: int | None = None
    best_area_pixels = float("inf")
    for line_index, line_text in enumerate(ocr_result.txts):
        line_normalized = _normalize_ocr_match_string(line_text)
        if substring_normalized not in line_normalized:
            continue
        box = ocr_result.boxes[line_index]
        xs = [float(p[0]) for p in box]
        ys = [float(p[1]) for p in box]
        area_pixels = (max(xs) - min(xs)) * (max(ys) - min(ys))
        if area_pixels < best_area_pixels:
            best_area_pixels = area_pixels
            best_line_index = line_index
    return best_line_index


def _first_logical_line_only_from_multiline_anchor_string(raw_anchor_string: str) -> str:
    """只保留第一逻辑行（首个换行前），用于列表页标题锚点：第二行起忽略。"""
    normalized_newlines_only = (raw_anchor_string or "").replace("\r\n", "\n").replace(
        "\r",
        "\n",
    )
    return normalized_newlines_only.split("\n", 1)[0].strip()


def ocr_find_text_center_allowing_note_title_wrapped_across_two_ocr_lines(
    image: str | Path | bytes | np.ndarray,
    target_text: str,
) -> tuple[int, int] | None:
    anchor_first_logical_line_only = _first_logical_line_only_from_multiline_anchor_string(
        target_text,
    )
    nested_recognizable_segment_string_list = (
        _split_anchor_raw_string_into_logical_line_then_recognizable_segment_nested_string_list(
            anchor_first_logical_line_only,
        )
    )
    if not nested_recognizable_segment_string_list:
        return None
    ocr = _get_rapid_ocr()
    result = ocr(_rapidocr_input_from_path_bytes_or_bgr_numpy(image))
    if result is None or not result.txts or result.boxes is None:
        return None
    logical_line_representative_center_xy_float_pair_list: list[tuple[float, float]] = []
    for recognizable_segment_raw_string_list_for_one_logical_line in nested_recognizable_segment_string_list:
        segment_center_xy_float_pair_list_for_one_logical_line: list[tuple[float, float]] = []
        for raw_segment_string in recognizable_segment_raw_string_list_for_one_logical_line:
            normalized_segment_string = _normalize_ocr_match_string(raw_segment_string)
            if not normalized_segment_string:
                continue
            matched_line_index = _ocr_smallest_area_line_index_where_normalized_line_contains_substring(
                result,
                normalized_segment_string,
            )
            if matched_line_index is None:
                return None
            segment_center_x, segment_center_y = _quad_center_xy(
                result.boxes[matched_line_index],
            )
            segment_center_xy_float_pair_list_for_one_logical_line.append(
                (segment_center_x, segment_center_y),
            )
        if not segment_center_xy_float_pair_list_for_one_logical_line:
            continue
        one_logical_line_representative_center_xy_float_pair = (
            _midpoint_float_xy_pair_from_first_and_last_center_in_ordered_float_pair_list(
                segment_center_xy_float_pair_list_for_one_logical_line,
            )
        )
        if one_logical_line_representative_center_xy_float_pair is not None:
            logical_line_representative_center_xy_float_pair_list.append(
                one_logical_line_representative_center_xy_float_pair,
            )
    if not logical_line_representative_center_xy_float_pair_list:
        return None
    final_anchor_center_xy_float_pair = (
        _midpoint_float_xy_pair_from_closest_pair_of_points_in_float_pair_list(
            logical_line_representative_center_xy_float_pair_list,
        )
    )
    if final_anchor_center_xy_float_pair is None:
        return None
    return (
        int(round(final_anchor_center_xy_float_pair[0])),
        int(round(final_anchor_center_xy_float_pair[1])),
    )


def ocr_find_text_center(
    image: str | Path | bytes | np.ndarray,
    target_text: str,
    *,
    prefer_highest_score: bool = True,
) -> tuple[int, int] | None:
    """
    对 ``image``（路径、``Path``、图像字节或 OpenCV BGR ``numpy`` 数组）跑 RapidOCR，找到与 ``target_text`` 最匹配的一行，
    返回该行四边形框的**中心点** ``(cx, cy)``（与图像像素坐标一致）。

    匹配规则：识别行经空白归一后，必须 **完整包含** ``target_text``（连续子串），
    不要求逐字与目标一致，但必须整段关键词都在同一 OCR 框对应的文字里。

    多行命中时：优先与目标 **完全一致** 的行，否则取识别串 **更短** 的行（更接近整块占位符），
    再按 ``prefer_highest_score`` 用置信度打破平局。
    """
    ocr = _get_rapid_ocr()
    result = ocr(_rapidocr_input_from_path_bytes_or_bgr_numpy(image))
    if result is None or not result.txts or result.boxes is None:
        return None

    target_norm = _normalize_ocr_match_string(target_text)
    if not target_norm:
        return None

    matches: list[tuple[int, float, int, int]] = []
    for i, txt in enumerate(result.txts):
        if not _ocr_line_contains_full_target(txt, target_text):
            continue
        line_norm = _normalize_ocr_match_string(txt)
        exact = 0 if line_norm == target_norm else 1
        line_len = len(line_norm)
        score = float(result.scores[i]) if result.scores and i < len(result.scores) else 0.0
        matches.append((i, score, exact, line_len))

    if not matches:
        return None

    if len(matches) > 1:
        if prefer_highest_score:
            matches.sort(key=lambda row: (row[2], row[3], -row[1]))
        else:
            matches.sort(key=lambda row: (row[2], row[3], row[1]))

    idx = matches[0][0]
    box = result.boxes[idx]
    cx, cy = _quad_center_xy(box)
    return int(round(cx)), int(round(cy))


def ocr_find_text_center_by_shortening_prefix_of_anchor_until_match(
    image: str | Path | bytes | np.ndarray,
    anchor_text: str,
    *,
    minimum_raw_character_count_inclusive: int = 6,
    prefer_highest_score: bool = True,
) -> tuple[int, int] | None:
    """
    在整段 ``anchor_text`` 无法命中时，从去掉末尾一字的前缀起，逐步缩短，
    直到某一前缀在单行 OCR 中可子串匹配（仍用 ``ocr_find_text_center`` 规则）。
    用于识别结果漏掉标题末尾一两字、或个别字与 DOM 不一致时仍能落到标题行。
    """
    stripped_anchor_text = (anchor_text or "").strip()
    character_count = len(stripped_anchor_text)
    if character_count <= minimum_raw_character_count_inclusive:
        return None
    for end_exclusive in range(
        character_count - 1,
        minimum_raw_character_count_inclusive - 1,
        -1,
    ):
        prefix_only_anchor_text = stripped_anchor_text[:end_exclusive]
        found_center_xy_integer_pair = ocr_find_text_center(
            image,
            prefix_only_anchor_text,
            prefer_highest_score=prefer_highest_score,
        )
        if found_center_xy_integer_pair is not None:
            return found_center_xy_integer_pair
    return None


def start(argv: list[str] | None = None) -> int:
    rest = argv if argv is not None else sys.argv[1:]
    if len(rest) < 2:
        print("用法: ocr-pos.py <目标文字> <图片路径>", file=sys.stderr)
        return 2

    target = (rest[0] or "").strip()
    image_path = Path(rest[1]).expanduser()
    if not target:
        print("目标文字不能为空。", file=sys.stderr)
        return 2
    if not image_path.is_file():
        print(f"图片不存在: {image_path}", file=sys.stderr)
        return 2

    pt = ocr_find_text_center(image_path, target)
    if pt is None:
        print(f"未在图中找到与「{target}」匹配的文字。", file=sys.stderr)
        return 1

    print(f"{pt[0]} {pt[1]}")
    return 0


def main() -> int:
    return start()


__all__ = [
    "REPO_ROOT",
    "OCR_RAPIDOCR_GLOBAL_MAX_SIDE_LEN",
    "OCR_RAPIDOCR_INFERENCE_ENGINE_NAME",
    "OCR_RAPIDOCR_ONNX_RUNTIME_USE_CUDA",
    "ocr_find_text_center",
    "ocr_find_text_center_allowing_note_title_wrapped_across_two_ocr_lines",
    "ocr_find_text_center_by_shortening_prefix_of_anchor_until_match",
    "ocr_smallest_line_center_for_dom_text",
    "start",
    "main",
]


if __name__ == "__main__":
    raise SystemExit(start())