project/tbsg/ai_helper_vision1.py

"""
AI 视觉助手1 — 列表页商品识别（最终版）
=============================================================================
流程：PP-OCRv6 识别（左右切割：物理裁掉左列商品图，秒级，生僻字强，坐标=真实屏幕像素）
     → GLM-5.2 分卡整理（V5提示词：卡片结构 + 残卡丢弃 + 推荐区忽略）
     → 几何校验（价格/店铺/标题逐卡反查回贴：宁丢卡不配错，与类目无关）
返回格式与 ai_helper_vision.VisionParser 一致，可直接替换。
独立测试：PyCharm 里直接 Run（配置在 __main__ 的 TEST_IMAGE / TEST_KEYWORD）
"""
import json
import os
import re
import sys
from typing import Optional

import requests

# ── 配置 ────────────────────────────────────────────────
GLM_URL = "https://ai-api.dfwy.tech/v1/messages"   # 中转站（Anthropic 格式端点）
GLM_MODEL = "glm-5.2"
GLM_KEY = "sk-Jzfd31hxASmK7nr9RyFfIEnn7dIA2e7zucfoovbOw4Hbo3kP"   # 中转站令牌

_PP_OCR = None   # PP-OCRv6 懒加载单例（模型初始化需数秒，进程内复用）

# ── 几何校验常量（按屏幕尺寸比例，类目无关，实测标定@1220x2712）──
GAP_RATIO = 0.08          # 价格→店铺最大垂直距离216px（同卡实测82~152px，跨卡≥577px）
TITLE_ABOVE_RATIO = 0.12  # 标题验证/回贴窗口（价格上方325px）
TITLE_BELOW_RATIO = 0.19  # 缝合卡检测窗口（价格下方515px）
RIGHT_COL_RATIO = 0.37    # 右列文字x中线阈值（过滤商品包装图噪声，包装字在更左侧）
MIN_TITLE_LEN = 6         # 标题块最低字符数（功效行/截断碎片≈4字，真标题≥6字）
# 促销/配送行特征（识别"店铺行"时排除用；满\d/减\d 带数字，避免误伤"满堂堂/减脂"类店名）
AID_PROMO_RE = re.compile(r"起送|满\d|配送|发货|日达|分钟|准时|红包|已售|月售|预估|减\d|新客|专享|原研")


def _get_ppocr():
    """懒加载 PP-OCRv6"""
    global _PP_OCR
    if _PP_OCR is None:
        from paddleocr import PaddleOCR
        _PP_OCR = PaddleOCR(
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
            enable_mkldnn=False,   # 绕过 Windows CPU oneDNN 推理 bug
        )
    return _PP_OCR


class VisionParser1:
    """PP-OCRv6 识别 + GLM 逐卡整理"""

    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or GLM_KEY

    # ── 主流程 ──────────────────────────────────────────

    def parse_shops(self, image_path: str, screen_size: tuple = (720, 1640),
                    keyword: str = "", crop_x: int = 0) -> tuple:
        """返回 (shops, status)：status = "ok" | "page_wrong"(页面不是药品列表) | "failed"(GLM失败)"""
        screen_w, screen_h = screen_size
        # 1. PP-OCRv6 识别（先物理裁掉左列商品图，只看右列文字区）
        blocks = self._ppocr_blocks(image_path, crop_x)
        print(f"[v1] ---------- PP-OCRv6 块(文字+坐标) ----------")
        for i, b in enumerate(blocks, 1):
            box = b["box"]
            print(f"  ocr[{i:02d}] box=[{box[0]},{box[1]}~{box[2]},{box[3]}] {b['text']}")
        if not blocks:
            print("[v1] PP-OCRv6 未识别到文字块（页面可能没加载出来）")
            return [], "failed"
        # 2. 全部文字块打包一次丢给 GLM：AI 自己划分卡片 + 提取字段（不做硬代码分组）
        shops, status = self._extract_all_cards(blocks, screen_w, screen_h, keyword)
        print(f"[v1] 最终 {len(shops)} 个店铺 (status={status})")
        return shops, status

    # ── PP-OCRv6 识别（左右切割：先物理裁掉左列商品图再识别）──

    def _ppocr_blocks(self, image_path: str, crop_x: int = 0) -> list:
        """PP-OCRv6 识别。crop_x>0 时先把左列商品图【物理切割】掉，
        只识别右列文字区（模型完全看不到左列干扰），坐标加回偏移还原。
        返回 [{text, box, bbox}]（box 为真实屏幕坐标 [x1,y1,x2,y2]）"""
        import numpy as np
        from PIL import Image
        ocr = _get_ppocr()
        img = Image.open(image_path).convert("RGB")
        W, H = img.size
        if crop_x and crop_x < W:
            seg = img.crop((crop_x, 0, W, H))   # 物理切割：只留右列
            xoff = crop_x
        else:
            seg = img
            xoff = 0
        blocks = []
        for res in ocr.predict(np.asarray(seg)):
            try:
                texts, boxes = res["rec_texts"], res["rec_boxes"]
            except Exception:
                j = res.json["res"]
                texts, boxes = j["rec_texts"], j["rec_boxes"]
            for t, bx in zip(texts, boxes):
                box = [int(v) for v in bx]
                box[0] += xoff
                box[2] += xoff
                x1, y1, x2, y2 = box
                if y2 < 80:
                    continue   # 状态栏
                blocks.append({"text": str(t), "box": box,
                               "bbox": [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]})
        return blocks

    # ── GLM 批量整理（全部文字块一次调用：AI自己分卡+提取）──

    def _extract_all_cards(self, blocks: list, screen_w: int, screen_h: int, keyword: str = "") -> tuple:
        """全部 OCR 文字块打包一次丢给 GLM，AI 自己划分卡片并提取字段。
        返回 (shops, status)：status = "ok" | "page_wrong"(页面不是药品列表) | "failed"(GLM失败)"""
        bs = sorted(blocks, key=lambda b: b["box"][1])
        lines = []
        for b in bs:
            box = b["box"]
            cx = (box[0] + box[2]) // 2
            cy = (box[1] + box[3]) // 2
            lines.append(f"[x={cx:4d}, y={cy:4d}] {b.get('text','')}")
        ocr_text = "\n".join(lines)
        system_prompt = self._cards_prompt()
        user_text = (
            f"屏幕OCR识别结果（共{len(bs)}个文字块）：\n\n{ocr_text}\n\n"
            f"当前搜索的商品是「{keyword}」。请划分商品卡片并逐卡提取字段。"
        )
        resp = self._call_llm(system_prompt, user_text)
        cleaned = (resp or "").strip()
        if cleaned.startswith("```"):
            ls = cleaned.split("\n")
            if ls and ls[0].startswith("```"):
                ls = ls[1:]
            if ls and ls[-1].strip() == "```":
                ls = ls[:-1]
            cleaned = "\n".join(ls).strip()
        data = None
        try:
            data = json.loads(cleaned)
        except json.JSONDecodeError:
            m = re.search(r'\[[\s\S]*\]', cleaned)
            if m:
                try:
                    data = json.loads(m.group(0))
                except json.JSONDecodeError:
                    print(f"[v1] GLM 返回JSON解析失败: {cleaned[:200]}")
        if not isinstance(data, list):
            print(f"[v1] GLM 返回不是数组，按失败处理")
            return [], "failed"
        # GLM判定页面不是药品列表（美食/外卖/首页等频道）→ 交给上层走恢复流程
        if len(data) == 1 and isinstance(data[0], dict) and data[0].get("page_wrong"):
            print(f"[v1] GLM 判定: 当前页面不是药品搜索列表（page_wrong）")
            return [], "page_wrong"
        shops = []
        for i, item in enumerate(data, 1):
            if not isinstance(item, dict):
                continue
            title = str(item.get("title", "") or "")
            price = str(item.get("price", "") or "")
            shop = str(item.get("shop", "") or "")
            sales = str(item.get("sales", "") or "")
            click = item.get("click") or item.get("coords") or []
            if not (shop and title and price):
                print(f"[v1] 卡{i}: [丢弃] 字段不完整: {shop} | {title[:20]} | {price}")
                continue
            if '￥' in title or '¥' in title:
                print(f"[v1] 卡{i}: [丢弃] 标题含价格符号（残卡）")
                continue
            if isinstance(click, list) and len(click) == 2:
                try:
                    cx, cy = int(click[0]), int(click[1])
                except (TypeError, ValueError):
                    continue
                if not (0 <= cx <= screen_w and 0 <= cy <= screen_h):
                    print(f"[v1] 卡{i}: [丢弃] 坐标越界: ({cx},{cy})")
                    continue
            else:
                print(f"[v1] 卡{i}: [丢弃] 无有效点击坐标")
                continue
            shop_y = int(item.get("shop_y", cy)) if str(item.get("shop_y", "")).lstrip('-').isdigit() else cy
            print(f"[v1] 卡{i}: [收录] {shop} | {title[:36]} | {price} | click=({cx},{cy}) | 已售={sales}")
            shops.append([shop, title, price, [cx, cy], shop_y, sales])
        # 几何校验：宁丢卡不配错（残卡泄漏/店铺错配/标题错位/形近字在此修正）
        shops = self._validate_cards(shops, bs, screen_w, screen_h)
        return shops, "ok"

    # ── 几何校验（GLM返回后逐卡验证、回贴、丢弃；规则全部基于布局几何，无领域词）──

    @staticmethod
    def _gnorm(s) -> str:
        return re.sub(r"\s+", "", str(s or "")).replace("（", "(").replace("）", ")").lower()

    @staticmethod
    def _gcy(b: dict) -> int:
        return (b["box"][1] + b["box"][3]) // 2

    @staticmethod
    def _gcx(b: dict) -> int:
        return (b["box"][0] + b["box"][2]) // 2

    @staticmethod
    def _glcs(a: str, b: str) -> int:
        """最长公共子串长度"""
        best = 0
        prev = [0] * (len(b) + 1)
        for ca in a:
            cur = [0]
            for j, cb in enumerate(b, 1):
                cur.append(prev[j - 1] + 1 if ca == cb else 0)
                best = max(best, cur[-1])
            prev = cur
        return best

    @classmethod
    def _price_block(cls, blocks: list, price: str, click_y) -> Optional[dict]:
        """R1: 找含该价格数字的￥块（取离click_y最近的）。找不到=编造价格"""
        m = re.search(r"[\d.]+", str(price or ""))
        if not m:
            return None
        pn = m.group(0).rstrip(".")
        cands = [b for b in blocks if ("¥" in b["text"] or "￥" in b["text"])
                 and re.search(r"(?<![\d.])" + re.escape(pn) + r"(?![\d])", cls._gnorm(b["text"]))]
        if not cands:
            return None
        ref = click_y[1] if isinstance(click_y, (list, tuple)) and len(click_y) == 2 else 10 ** 9
        return min(cands, key=lambda b: abs(cls._gcy(b) - ref))

    @classmethod
    def _shop_text(cls, text: str) -> str:
        """店铺最终文本 = OCR块文本取括号前"""
        return re.split(r"[(（]", re.sub(r"\s+", "", str(text or "")))[0]

    @classmethod
    def _shop_block(cls, blocks: list, shop: str, y_from: int, y_to: int) -> Optional[dict]:
        """R2: [y_from,y_to)窗口内找店铺名块。
        容错匹配：包含关系 或 括号前店名等长仅差1字（防 沪/沚/滪 形近字）"""
        ns = cls._gnorm(shop)
        if len(ns) < 3:
            return None
        best, best_score = None, 0
        for b in blocks:
            y = cls._gcy(b)
            if not (y_from <= y < y_to):
                continue
            ok = False
            for tt in (cls._gnorm(b["text"]), cls._gnorm(cls._shop_text(b["text"]))):
                if len(tt) < 3:
                    continue
                if ns in tt or tt in ns:
                    ok = True
                    break
                if len(tt) == len(ns) and len(ns) >= 4 and sum(a != c for a, c in zip(tt, ns)) == 1:
                    ok = True
                    break
            if ok:
                score = min(len(ns), len(cls._gnorm(b["text"])))
                if score > best_score:
                    best, best_score = b, score
        return best

    @classmethod
    def _shop_row_repair(cls, blocks: list, y_from: int, y_to: int) -> Optional[dict]:
        """R2修复: GLM把店名整列配错位时（典型=顶部残卡店名强配下一卡，引发级联），
        按几何重新定位本卡的店铺行。店铺行特征：非促销、非价格、长度≥4，
        且同一行(±40px)有 发货/起送/日达/分钟 等伴随块（标题行从来没有伴随块）。"""
        for b in sorted((x for x in blocks if y_from <= cls._gcy(x) < y_to), key=cls._gcy):
            t = cls._gnorm(b["text"])
            if len(t) < 4 or "¥" in b["text"] or "￥" in b["text"] or AID_PROMO_RE.search(t):
                continue
            by = cls._gcy(b)
            has_aid = any(abs(cls._gcy(a) - by) <= 40 and re.search(r"发货|起送|日达|分钟|准时|配送", cls._gnorm(a["text"]))
                          for a in blocks)
            if has_aid:
                return b
        return None

    @classmethod
    def _title_blocks(cls, blocks: list, title: str, y_from: int, y_to: int, right_cx: int) -> list:
        """窗口内右列与标题匹配的块（包含关系 或 公共子串≥5），按匹配长度降序"""
        nt = cls._gnorm(title)
        if len(nt) < 4:
            return []
        out = []
        for b in blocks:
            y = cls._gcy(b)
            if not (y_from <= y < y_to) or cls._gcx(b) < right_cx:
                continue
            t = cls._gnorm(b["text"])
            if len(t) < 4:
                continue
            if nt in t or t in nt or cls._glcs(t, nt) >= 5:
                out.append((min(len(t), len(nt)), b))
        out.sort(key=lambda x: -x[0])
        return [b for _, b in out]

    @classmethod
    def _validate_cards(cls, cards: list, blocks: list, screen_w: int, screen_h: int) -> list:
        """几何校验：宁丢卡不配错。cards元素 = [shop,title,price,[cx,cy],shop_y,sales]
        R1 价格块可反查（防编造价格/坐标错卡）
        R2 店铺块在价格下方窗口内可反查（防店铺配到别的卡）
        R4 一个店铺块只属于一张卡（防一店名连配多卡）
        R5 标题可在价格上方验证到（≥6字；只在其下方=缝合卡，无法验证=碎片拼接）
        R6 店名/标题回贴为OCR原文（修正形近字/抄写走样，保留GLM多行拼接的标题）"""
        max_gap = max(int(screen_h * GAP_RATIO), 80)
        t_above = int(screen_h * TITLE_ABOVE_RATIO)
        t_below = int(screen_h * TITLE_BELOW_RATIO)
        right_cx = int(screen_w * RIGHT_COL_RATIO)
        kept, used_shop = [], []
        for s in cards:
            shop, title, price, click = s[0], s[1], s[2], s[3]
            tag = f"{shop}|{price}|{title[:14]}"
            pb = cls._price_block(blocks, price, click)                    # R1
            if pb is None:
                print(f"[v1] [校验丢弃·R1价格块不存在] {tag}")
                continue
            py = cls._gcy(pb)
            sb = cls._shop_block(blocks, shop, py + 5, py + max_gap)       # R2
            repaired = False
            if sb is None:
                # 店名在价格下方找不到 → 多半是GLM把店名整列配错位（残卡引发级联）
                # 按几何找本卡真正的店铺行修复，救回整批卡片而不是全部丢弃
                sb = cls._shop_row_repair(blocks, py + 5, py + max_gap)
                if sb is None:
                    print(f"[v1] [校验丢弃·R2价格下方无此店] {tag}")
                    continue
                repaired = True
            sy = cls._gcy(sb)
            if any(abs(sy - u) < 25 for u in used_shop):                   # R4
                print(f"[v1] [校验丢弃·R4店铺块复用] {tag}")
                continue
            above = [b for b in cls._title_blocks(blocks, title, py - t_above, py, right_cx)
                     if len(cls._gnorm(b["text"])) >= MIN_TITLE_LEN]
            below = cls._title_blocks(blocks, title, py + 5, py + t_below, right_cx)
            if not above:                                                  # R5
                reason = "标题偷自下方(缝合卡)" if below else "标题无法验证"
                print(f"[v1] [校验丢弃·R5{reason}] {tag}")
                continue
            # R6 回贴：GLM标题包含OCR块→保留GLM（多行拼接的规格）；否则用OCR块（修正错位/形近字）
            t1 = re.sub(r"\s+", "", str(title))
            t2 = re.sub(r"\s+", "", str(above[0]["text"]))
            new_title = t1 if len(t1) > len(t2) and cls._gnorm(t2) in cls._gnorm(t1) else t2
            s[0] = cls._shop_text(sb["text"])
            s[1] = new_title
            s[3] = [cls._gcx(pb), py]      # 点击坐标贴到真实价格块中心
            s[4] = sy                      # shop_y 贴到真实店铺块
            used_shop.append(sy)
            fix_note = " [R2修复·店名错位]" if repaired else ""
            print(f"[v1] [校验通过]{fix_note} {s[0]} | {s[1][:36]} | {price} | click=({s[3][0]},{s[3][1]})")
            kept.append(s)
        return kept

    def _cards_prompt(self) -> str:
        """整页整理提示词：AI 自己划分卡片 + 提取字段（不做硬代码分组）
        V5：V0 基础上打 3 个补丁——顶部孤儿店铺名规则 / 推荐区忽略 / 店铺名定义去类目化"""
        return """你是电商列表页 OCR 解析助手。收到一张饿了么闪购搜索结果页的 OCR 文字（每行格式 [x=横坐标, y=纵坐标] 文字，y 越大越靠屏幕下方）。

【任务】自己把文字块划分成商品卡片，再逐卡提取字段。

【卡片划分方法】
- 卡片竖向排列：同一卡片的文字块 y 坐标相近；两个文字块的 y 间隔超过约 80px 就属于不同卡片
- 页头（如"30分钟达""适用人群""筛选"）不是商品卡片，忽略
- 分隔说明行（如"搜索结果较少，为你推荐相关店铺"）之后的推荐内容不是搜索结果卡片，全部忽略：
  推荐店铺（店名/评分/月售/起送/距离）以及其中夹着的小商品瓦片（商品名+价格），
  严禁把推荐店铺名和瓦片商品配对成卡片输出
- 每张卡片结构固定，从上到下：
  1. 商品标题（品牌标签如[理洫王]、[999] + 药名 + 规格，可能多行）
  2. （中间可能夹杂：处方药、OTC 等标签行——忽略）
  3. 价格（￥或¥开头，只有一行）
  4. （下面可能夹杂：满减、起送、指定商品、配送费、快递发货、1-3日达等文字——忽略）
  5. 店铺名（价格下方的商家/店铺名称）

【提取方法】（逐卡独立提取，不得把别的卡片的文字拼进来）
- price = 本卡内以￥或¥开头的块。若一块里有两个价格（如"¥33¥38"），取第一个（主价），第二个是划线原价
- title = 价格【上方】的本卡商品名文字，多行完整拼接
  （规格直接相连，如「0.33g*2」+「4粒/盒」拼成「0.33g*24粒/盒」）；跳过"处方药"等标签行；保留开头的品牌标签
- shop = 价格【下方】的商家名，含"（"只取"（"前；满减/起送/快递发货/日达这些行不是店铺名
- click = [x, y]：直接取【价格块】的坐标（价格在商品信息中部，点价格进商品最稳，不会点错卡）
- shop_y = 店铺名块的 y
- sales = "已售"后面的数字，没有留空""

【残卡判断——标题、价格、店铺名三者缺一不可】
完整卡片必须同时具备：标题（价格上方）+ 价格 + 店铺名（价格下方）。
屏幕顶部/底部被截断的卡（缺标题或缺店铺名）直接不输出，不要强行拼凑。
特别注意屏幕顶部的残卡：如果一段内容直接以价格开头（它上方没有本卡的标题），
说明标题已滚出屏幕——这个价格和它紧下方的店铺名属于同一张被截断的卡，整卡丢弃；
严禁把这个店铺名（或这个价格）配给屏幕上其他任何一张完整卡片。
同样，屏幕底部最后一段若只有标题和价格、下面没有店铺名，也整卡丢弃。
其余正常卡片都正常提取，不要因为文字杂乱、有促销行或距离问题就丢卡。

【页面判断——先看页面对不对】
如果整个页面不是商品搜索结果列表（比如首页、美食/外卖商家频道、商品详情页等），
不要硬凑卡片，返回 [{"page_wrong": true}]。
特别提醒：采集中退回到错误页面时，页面上可能是餐饮等其他类目的商家，
它们不是当前搜索的商品，绝不能当卡片提取。

【约束】
- 所有字段严格取自 OCR 原文，坐标用 OCR 里的真实坐标，不得编造
- 输出纯 JSON 数组（字符串用双引号），不要任何多余文字：
[{"title":"商品标题","price":"价格","shop":"店铺名","click":[x,y],"shop_y":店铺名y,"sales":"月售"}, ...]"""

    def _call_llm(self, system_prompt: str, user_text: str, timeout: int = 300) -> str:
        """调中转站 GLM（Anthropic 格式端点，关思考防吃光 token）"""
        headers = {"x-api-key": GLM_KEY, "anthropic-version": "2023-06-01",
                   "Content-Type": "application/json"}
        payload = {"model": GLM_MODEL, "max_tokens": 8192,
                   "system": system_prompt, "thinking": {"type": "disabled"},
                   "messages": [{"role": "user", "content": user_text}]}
        last_err = None
        for attempt in range(3):
            try:
                r = requests.post(GLM_URL, headers=headers, json=payload, timeout=timeout)
                if r.status_code != 200:
                    print(f"[v1] GLM HTTP {r.status_code}: {r.text[:200]}")
                    last_err = Exception(f"HTTP {r.status_code}")
                    import time as _t
                    _t.sleep(2); continue
                data = r.json()
                if "error" in data:
                    err = data["error"].get("message", str(data["error"])) if isinstance(data["error"], dict) else str(data["error"])
                    print(f"[v1] GLM error: {err[:200]}")
                    last_err = Exception(err)
                    import time as _t
                    _t.sleep(2); continue
                texts = [c.get("text", "") for c in data.get("content", []) if c.get("type") == "text"]
                return "".join(texts)
            except Exception as e:
                last_err = e
                print(f"[v1] GLM 调用失败(第{attempt+1}次)({e})，重试...")
                import time as _t
                _t.sleep(2)
        print(f"[v1] GLM 3次均失败: {last_err}")
        return ""


# 别名：方便 main1.py 用 `from ai_helper_vision1 import VisionParser` 直接替换
VisionParser = VisionParser1


# ── 独立测试 ───────────────────────────────────────────
if __name__ == "__main__":
    # ═══════ 测试配置：改这里就行，不用命令行传参 ═══════
    TEST_IMAGE = r""  # 要测的图片完整路径；留空 = 自动找最新 step3_b0.png
    TEST_KEYWORD = "血塞通软胶囊"  # 搜索关键词
    TEST_CROP_X = 287  # 左列分界（列表页一般 287~289）；0 = 不过滤
    # ═══════════════════════════════════

    kw = TEST_KEYWORD

    import glob
    import re as _re
    _proj_root = os.path.dirname(os.path.abspath(__file__))
    _shot_root = os.path.join(_proj_root, "screenshots")
    _all = []
    if os.path.isdir(_shot_root):
        for p in glob.glob(os.path.join(_shot_root, "**", "step3", "step3_b0.png"), recursive=True):
            if _re.match(r'^step3_b\d+\.png$', os.path.basename(p)):
                _all.append(p)
    _all.sort(key=os.path.getmtime, reverse=True)

    if TEST_IMAGE.strip():
        _all = [TEST_IMAGE.strip()]

    if not _all:
        print(f"没找到截图（在 {_shot_root}）")
        sys.exit(1)

    print(f"找到 {len(_all)} 张截图，关键词: {kw}\n")

    for img_path in _all:
        print("=" * 60)
        print(f"图片: {img_path}")
        if not os.path.exists(img_path):
            print("文件不存在，跳过"); continue
        try:
            from PIL import Image
            w, h = Image.open(img_path).size
        except ImportError:
            w, h = 720, 1640
        try:
            parser = VisionParser1()
            shops, status = parser.parse_shops(img_path, screen_size=(w, h), keyword=kw, crop_x=TEST_CROP_X)
            print(f"  -> 识别 {len(shops)} 个店铺:")
            for s in shops:
                print(f"     {s[0]} | {s[1][:36]} | {s[2]} | click={s[3]} | 已售={s[5]}")
        except Exception as e:
            print(f"  异常: {e}")
        print()