project/tbsg/ai_helper_vision2.py

# -*- coding: utf-8 -*-
"""
ai_helper_vision2 — 起送行锚定列表解析(纯规则, 零AI调用)
=============================================================================
替换 ai_helper_vision1 的 PP-OCR+GLM 交叉验证分卡。

原理(2026-09-20 与用户真机联调确定):
  每张商品卡片的底部必然是一行「起送¥…」(与"1-3日达"同行)
  → 起送行 = 卡片底部锚点, 相邻两条起送行之间 = 一张卡的完整内容
  → 卡片数 = 起送行数, 结构性防漏卡/防错位
  残卡: 首区间高度 < 中位60%, 或 首区间文字不含采集关键词 → 裁掉不采
  底部: 最后一条起送行以下留给下一屏, 不产生残卡
  耗尽横幅: 「搜索结果较少,为你推荐相关店铺」以下 = 店铺推荐流, 不采集
  店铺名: 只取括号前部分(去「(快递发…)」等后缀)
  字段: 标题/价格/已售/店铺名 全部规则提取, 失败只影响字段不影响分卡

接口与 ai_helper_vision1 完全兼容:
  parse_shops(image_path, screen_size=(w,h), keyword='', crop_x=0) -> (shops, status)
  shops 元素: [店名, 标题, 价格, [click_x, click_y], 店铺行y, 已售]
  status: "ok" | "page_wrong"(无起送行,非列表页) | "failed"(OCR失败)

独立测试: python ai_helper_vision2.py <截图路径> [更多路径] [--kw 采集关键词]
调试图自动保存: swipe_test_preview/vision2_debug/
"""
import os
import re

import cv2
import numpy as np

try:                                    # main1 环境(steps已在sys.path)
    from ocr import OCR as _OCRClass
except ImportError:                     # 独立运行
    from steps.ocr import OCR as _OCRClass

OUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                       "swipe_test_preview", "vision2_debug")
os.makedirs(OUT_DIR, exist_ok=True)

HAS_Y = lambda t: ('¥' in t or '￥' in t)                    # 半角/全角人民币符号
SHOP_RE = re.compile(r'药房|药店|医药|健康|自营|旗舰|连锁|快递发')
BADGE = ('快递发货', '快递')                                  # 纯徽标块不算店铺名
TAIL_RE = re.compile(r'^起送|日达$')                          # 卡片底部行(起送/日达)
SPEC_RE = re.compile(r'\d|盒|丸|支|瓶|片|粒|袋|贴|枚')         # 标题续行特征(含规格)
EFF_RE = re.compile(r'湿疹|皮炎|瘙痒|疼痛|止咳|祛痰|消炎|杀菌|通络|活血'
                    r'|止痛|消肿|鼻炎|感冒|咳嗽|祛风|燥湿|止痒|解毒|清热|上火'
                    r'|手足癣|体股癣|花斑癣|神经性|异位性|脓疱|头皮脂溢')

_OCR_SINGLE = None


def _get_ocr():
    """RapidOCR 单例"""
    global _OCR_SINGLE
    if _OCR_SINGLE is None:
        _OCR_SINGLE = _OCRClass()
    return _OCR_SINGLE


class VisionParser2:
    """起送行锚定列表解析器(接口兼容 VisionParser1)"""

    def __init__(self, api_key: str = None):
        self.api_key = api_key      # 兼容保留, 本类不使用

    # ── 主入口(接口兼容 VisionParser1.parse_shops) ──
    def parse_shops(self, image_path: str, screen_size: tuple = (720, 1640),
                    keyword: str = "", crop_x: int = 0) -> tuple:
        try:
            o = _get_ocr()
            raw = o.recognize(image_path, detail="all")
            if not raw:
                print("[v2] OCR无结果")
                return [], "failed"
            H = max((r["box"][3] for r in raw), default=0) or screen_size[1]

            # 关键词: 未传时从顶部搜索框取最长块
            # (排除状态栏时钟 及 顶部chips/横幅词, 否则会把"8:38"当关键词)
            keyword = (keyword or "").strip()
            if not keyword:
                top_blocks = [r["text"] for r in raw
                              if (r["box"][1] + r["box"][3]) // 2 < H * 0.085
                              and r["box"][1] >= 60                      # 跳过状态栏
                              and not re.match(r'^\d{1,2}:\d{2}$', r["text"])  # 跳过时钟
                              and r["text"].strip() not in
                              ('外卖红包', '医保支付', '快递', '阿里自营', '30分钟达',
                               '筛选', '签到', '搜索')]
                if top_blocks:
                    keyword = max(top_blocks, key=len).strip()

            # 右列过滤(裁剪坐标: 优先调用方给的 crop_x, 否则用¥块左缘中位数自检)
            if not crop_x or crop_x <= 0:
                ysb = [r["box"][0] for r in raw if HAS_Y(r.get("text", ""))]
                crop_x = int(np.median(ysb)) - 15 if ysb else int(screen_size[0] * 0.4)
            blocks = sorted([((r["box"][1] + r["box"][3]) // 2, r["text"],
                              r["box"][3], (r["box"][0] + r["box"][2]) // 2)
                             for r in raw
                             if (r["box"][0] + r["box"][2]) // 2 >= crop_x])

            # ── 底部边界: 同一水平内出现「两个及以上」卡片尾部标记:
            #    ①起送¥ ②免配送费/配送¥ ③1-3日达 (②③可与①合并在同一块内,
            #    如"起送￥20配送￥5" = ①+② 两个标记) ──
            bounds, i = [], 0
            while i < len(blocks):
                j = i
                while j < len(blocks) and blocks[j][0] - blocks[i][0] <= 16:
                    j += 1
                grp = blocks[i:j]
                texts = [t for _, t, _, _ in grp]
                n = sum([any('起送' in t and HAS_Y(t) for t in texts),
                         any('免配送费' in t or '配送¥' in t or '配送￥' in t for t in texts),
                         any('日达' in t for t in texts)])
                if n >= 2 and grp[0][0] > H * 0.12:
                    bounds.append({'cy': grp[0][0], 'line': max(g[2] for g in grp) + 3})
                i = j
            if not bounds:
                print("[v2] 未检出起送行 → 非列表页")
                self._debug_save(image_path, [], [])
                return [], "page_wrong"

            # ── 顶线: 快递/阿里自营/筛选 chips 行(文字下方) ──
            top, top_line = 0, 0
            for a in range(len(blocks)):
                grp = [g for g in blocks if abs(g[0] - blocks[a][0]) <= 16 and g[0] < H * 0.18]
                texts = [t for _, t, _, _ in grp]
                n = sum([any(t.strip() == '快递' for t in texts),
                         any('阿里自营' in t for t in texts),
                         any('筛选' in t for t in texts)])
                if n >= 2:
                    top, top_line = grp[0][0], max(g[2] for g in grp) + 3
                    break

            # ── 区间划分与分类 ──
            # 耗尽横幅「搜索结果较少,为你推荐相关店铺」以下 = 店铺推荐流, 不采集
            banner_top = 0
            for r in raw:
                t = r.get('text', '')
                if ('搜索结果较少' in t) or ('为你推荐相关店铺' in t):
                    banner_top = r['box'][1] if not banner_top else min(banner_top, r['box'][1])

            cys = [b['cy'] for b in bounds]
            intervals = [(top_line, cys[0])] + [(cys[k], cys[k + 1]) for k in range(len(cys) - 1)]
            intervals.append((cys[-1], H))
            heights = [b - a for a, b in intervals[:-1]] or [1]
            med = sorted(heights)[len(heights) // 2]

            shops, marks = [], []
            for k, (a, b) in enumerate(intervals):
                ib = [(cy, t, cx) for cy, t, y2, cx in blocks if a < cy <= b]
                price_t = next((t for cy, t, cx in reversed(ib)
                                if HAS_Y(t) and not t.startswith('起送')), '')
                last = (k == len(intervals) - 1)
                if last:
                    marks.append((k + 1, 'bottom_partial', a, b, (160, 160, 160)))
                    continue
                if banner_top and (a >= banner_top - 5 or a < banner_top < b):
                    marks.append((k + 1, 'below_banner', a, b, (200, 150, 80)))
                    continue
                if (b - a) < 0.6 * med:
                    marks.append((k + 1, 'residual', a, b, (0, 0, 255)))
                    continue
                kw_hit = bool(keyword) and any(keyword[:2] in t or keyword[:4] in t
                                               for _, t, _ in ib)
                if a == top_line and keyword and not kw_hit:
                    marks.append((k + 1, 'residual', a, b, (0, 0, 255)))
                    continue
                # ── 字段提取 ──
                price_m = re.search(r'[¥￥]\s*([0-9]+(?:\.[0-9]+)?)', price_t)
                price = f"¥{price_m.group(1)}" if price_m else price_t[:8]
                sold_m = re.search(r'已售(\d+\+?)', price_t)
                sold = sold_m.group(1) if sold_m else ''
                title = self._build_title(ib)
                shop = self._extract_shop(ib, title)
                click = self._price_center(ib)
                shops.append([shop, title, price, click, bounds[k]['line'], sold])
                marks.append((k + 1, 'ok', a, b, (0, 200, 0)))

            n_ok = sum(1 for m in marks if m[1] == 'ok')
            n_res = sum(1 for m in marks if m[1] == 'residual')
            n_bot = sum(1 for m in marks if m[1] == 'bottom_partial')
            n_ban = sum(1 for m in marks if m[1] == 'below_banner')
            print(f"[v2] 起送行{len(bounds)} → 完整{n_ok} 残卡{n_res} 底部{n_bot} 横幅下{n_ban} "
                  f"(标注图: vision2_debug/)")
            self._debug_save(image_path, marks, bounds)
            return shops, "ok"
        except Exception as e:
            print(f"[v2] 解析异常: {e}")
            import traceback
            traceback.print_exc()
            return [], "failed"

    # ── 标题组装: 首块 + 规格续行(跳过功效行) ──
    def _build_title(self, ib):
        content = [t for cy, t, cx in ib if not TAIL_RE.search(t) and not HAS_Y(t)]
        if not content:
            return ''
        title = content[0].strip()
        if len(content) >= 2:
            nxt = content[1].strip()
            m = SPEC_RE.search(nxt)
            if len(nxt) <= 14 and m and not EFF_RE.search(re.sub(r'\d|盒|丸|支|瓶|片|粒|袋|贴|枚', '', nxt)):
                title += nxt
        return title[:60]

    # ── 店铺名: 数据字段, 宽松提取, 失败不影响分卡 ──
    def _extract_shop(self, ib, title):
        shop = next((t for cy, t, cx in reversed(ib)
                     if SHOP_RE.search(t) and t.strip() not in BADGE
                     and not (len(t.strip()) <= 5 and '快递' in t)), '')
        if shop:
            shop = re.split(r'快递发货', shop)[0]           # 剥掉混入的「快递发货」徽标
            shop = re.split(r'[（(]', shop)[0].strip()       # 只要括号前的店名(去「(快递发…)」后缀)
            shop = re.sub(r'[.．…]+$', '', shop).strip()    # 去尾部省略号
        if shop:
            return shop[:24]
        cand = [t for cy, t, cx in ib
                if t != title and not HAS_Y(t) and not TAIL_RE.search(t)
                and len(t.strip()) >= 4]
        if cand:
            return re.split(r'[（(]', max(cand, key=len))[0].strip()[:20]
        return ''

    # ── 价格块中心(click) ──
    def _price_center(self, ib):
        cand = [(cy, cx) for cy, t, cx in ib if HAS_Y(t) and not t.startswith('起送')]
        if not cand:
            return [0, 0]
        cy, cx = cand[-1]
        return [int(cx), int(cy)]

    # ── 调试图: 起送行绿线+各区间标注; 首区间为残卡时额外输出裁掉残卡的最终图 ──
    def _debug_save(self, image_path, marks, bounds):
        try:
            img = cv2.imread(image_path)
            if img is None:
                return
            for (n, kind, a, b, color) in marks:
                cv2.putText(img, f'#{n} {kind}', (8, min(a + 30, img.shape[0] - 10)),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
            for b in bounds:
                cv2.line(img, (0, b['line']), (img.shape[1], b['line']), (0, 200, 0), 2)
            base = os.path.basename(image_path)
            cv2.imwrite(os.path.join(OUT_DIR, f'{base}_v2.png'), img)
            # 最终图: 只留完整卡区域 — 首区间残卡→从第一条边界线裁顶,
            # 末区间底部不完整→裁到最后一条边界线为止
            crop_top = bounds[0]['line'] if marks and marks[0][1] == 'residual' else 0
            crop_bottom = bounds[-1]['line'] if marks and marks[-1][1] == 'bottom_partial' \
                else img.shape[0]
            final = img[crop_top:crop_bottom, :]
            if final.size:
                cv2.imwrite(os.path.join(OUT_DIR, f'{base}_v2_final.png'), final)
        except Exception:
            pass


VisionParser = VisionParser2


if __name__ == '__main__':
    import sys
    args = [a for a in sys.argv[1:] if a != '--kw']
    kw = sys.argv[sys.argv.index('--kw') + 1] if '--kw' in sys.argv else ''
    if not args:
        print('用法: python ai_helper_vision2.py <截图路径> [更多路径] [--kw 采集关键词]')
        sys.exit(0)
    p = VisionParser2()
    for f in args:
        print(f'\n===== {os.path.basename(f)} (关键词: {kw or "自动从搜索框提取"}) =====')
        shops, status = p.parse_shops(f, keyword=kw)
        print(f'status={status}, 卡片{len(shops)}张:')
        for s in shops:
            print(f'  {s[0]} | {str(s[1])[:36]} | {s[2]} | 已售={s[5]} | click={s[3]}')