aiData/Apps/AiTeJiYiChong/ReadImageKit.py

import numpy as np
from PIL import Image
import os
import asyncio
import hashlib
import json
import aiohttp
import logging
import base64
import cv2
import uuid
from openai import AsyncOpenAI, BadRequestError
from Config.Config import (
    ALY_LLM_API_KEY, VL_MODEL_NAME, VL_MODEL_NAME_AD, TEMP_IMAGE_DIR
)
from Apps.AiTeJiYiChong.Config.Setting import (
    SAFE_EXCLUDE_RATIO, FALLBACK_WIDTH, FALLBACK_HEIGHT,
    BOTTOM_SAFE_EXCLUDE_RATIO
)
from Apps.AiTeJiYiChong import Kit
from Apps.AiTeJiYiChong.Kit import (
    read_image, save_image, setup_logger
)

logger = setup_logger("ReadImageKit")

# 场站过滤黑名单
STATION_EXCLUDED_TITLES = {
    "推荐站点", "距离优先", "价格最低", "空闲较多", "筛选", "地图", "列表",
    "我的订单", "收藏站点", "优惠券", "站内搜索"
}

STATION_BLACKLIST_KEYWORDS = [
    "广告", "点击进入", "立即参与", "咨询电话", "活动详情", "导航到", "停车说明"
]


class ReadImageKit:
    _client = AsyncOpenAI(
        api_key=ALY_LLM_API_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 通用回退设备信息，仅在无法动态获取设备信息时使用
    _FALLBACK_DEVICE_INFO = {
        "displayWidth": FALLBACK_WIDTH,
        "displayHeight": FALLBACK_HEIGHT,
        "productName": "generic"
    }

    @staticmethod
    def _extract_json(text: str) -> str:
        if not text:
            return "[]"

        cleaned = text.strip()
        if "```" in cleaned:
            lines = []
            for line in cleaned.splitlines():
                if line.strip().startswith("```"):
                    continue
                lines.append(line)
            cleaned = "\n".join(lines).strip()

        decoder = json.JSONDecoder()

        pos = 0
        while pos < len(cleaned):
            idx_dict = cleaned.find("{", pos)
            idx_list = cleaned.find("[", pos)

            candidates = [i for i in (idx_dict, idx_list) if i != -1]
            if not candidates:
                break

            start = min(candidates)
            snippet = cleaned[start:]
            try:
                _, end = decoder.raw_decode(snippet)
                return snippet[:end]
            except json.JSONDecodeError:
                pos = start + 1
                continue

        return "[]"

    _prompt_detail = (
        "仅输出JSON对象（不含任何说明文字），识别充电站详情图片中的以下信息：\n"
        "1. station_name: 场站名称；\n"
        "2. address: 场站完整地址（通常在定位图标旁）。\n"
        "\n"
        "特别说明：\n"
        "- 严禁输出 Markdown 代码块标签，严格返回纯JSON对象。"
    )

    @classmethod
    async def detect_ad_popup(cls, image_path: str, device_info=None):
        """
        使用 VL 模型检测是否存在弹窗广告，并返回关闭按钮坐标
        """
        if not os.path.exists(image_path):
            return None

        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
            img_input = {"url": f"data:image/jpeg;base64,{encoded_image}"}

        # 构建提示词
        prompt = (
            "请仔细检查这张图片中是否存在**弹窗广告**或**悬浮广告**。\n"
            "广告可能有以下几种形式：\n"
            "1. **屏幕中央的大型弹窗广告**：通常遮挡了页面内容，内容多为优惠券、活动推广、或“免费停车”类提示。\n"
            "2. **悬浮广告**：一个小型的可点击图标，带有明显的关闭按钮。\n"
            "3. **底部的横幅广告**：带有明显的关闭按钮。\n\n"
            "请返回关闭按钮的中心坐标。\n\n"
            "请以纯 JSON 格式输出：\n"
            "{\n"
            "  \"has_ad\": true/false,\n"
            "  \"ad_type\": \"center\" | \"bottom\" | \"other\",\n"
            "  \"close_point\": [x, y]  // 0-1000 归一化坐标\n"
            "}\n"
            "如果没有广告，请返回 `{\"has_ad\": false}`。"
        )

        try:
            resp = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME, 
                timeout=30,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": img_input},
                            {"type": "text", "text": prompt},
                        ],
                    }
                ],
            )
            content = resp.choices[0].message.content.strip()
            raw_json = cls._extract_json(content)
            result = json.loads(raw_json)

            if result.get("has_ad") and result.get("close_point"):
                norm_point = result["close_point"]

                # 转换坐标
                if len(norm_point) == 2:
                    if not device_info:
                        device_info = cls._FALLBACK_DEVICE_INFO
                    
                    w = device_info.get("displayWidth", FALLBACK_WIDTH)
                    h = device_info.get("displayHeight", FALLBACK_HEIGHT)
                    
                    x = int(norm_point[0] / 1000 * w)
                    y = int(norm_point[1] / 1000 * h)
                    
                    return {
                        "x": x, 
                        "y": y,
                        "ad_type": result.get("ad_type", "unknown")
                    }
            
            return None

        except Exception as e:
            logger.error(f"Failed to detect ad popup: {e}")
            return None

    @staticmethod
    def _to_minutes(t_str: str) -> int:
        """HH:MM -> 分钟数"""
        if not t_str or ":" not in t_str:
            return 0
        try:
            h, m = map(int, t_str.split(":"))
            return h * 60 + m
        except:
            return 0

    @staticmethod
    def _fmt(t: int) -> str:
        """分钟数 -> HH:MM"""
        h = t // 60
        m = t % 60
        return f"{h:02d}:{m:02d}"

    @staticmethod
    def expand_schedule_to_24h(rows: list) -> list:
        """
        将时段列表规整为全天24个整点小时段
        """
        # 预处理：转换为分钟区间
        intervals = []
        for r in rows:
            s = ReadImageKit._to_minutes(r.get("start"))
            e = ReadImageKit._to_minutes(r.get("end"))
            if e <= s and e != 0: # 处理 00:00-00:00 这种可能
                if e == 0: e = 1440
                else: continue
            if e == 0 and s > 0: e = 1440
            
            s = max(0, s)
            e = min(1440, e)
            intervals.append({
                "s": s, "e": e,
                "price": r.get("price")
            })
        
        # 排序
        intervals.sort(key=lambda x: (x["s"], x["e"]))

        result = []
        for h in range(24):
            hs = h * 60
            he = (h + 1) * 60
            best_price = None
            
            # 找到覆盖当前小时段的价格
            for it in intervals:
                # 计算重叠
                overlap_s = max(hs, it["s"])
                overlap_e = min(he, it["e"])
                if overlap_e > overlap_s:
                    best_price = it["price"]
                    break # 简单处理，取第一个覆盖的
            
            result.append({
                "start": ReadImageKit._fmt(hs),
                "end": ReadImageKit._fmt(he),
                "price": best_price
            })
        return result

    _prompt_price_detail = (
        "仅输出JSON对象（不含任何说明文字），识别充电价格详情图片中的分时段电价信息。图片中通常包含“充电时段”、“单价”、“电费”、“服务费”等列。\n"
        "请识别出所有的时段行，返回一个列表，每个元素包含：\n"
        "1. start: 开始时间 (HH:MM)；\n"
        "2. end: 结束时间 (HH:MM)；\n"
        "3. price: 总单价 (数字，如 0.6100)；\n"
        "4. ele_fee: 电费 (数字)；\n"
        "5. ser_fee: 服务费 (数字)。\n"
        "\n"
        "注意：\n"
        "- 只需识别“直流桩”或当前选中的桩型下的数据。\n"
        "- 严禁输出 Markdown 代码块标签，严格返回纯JSON对象。"
    )

    @classmethod
    async def get_price_detail_from_image(cls, image_path: str):
        """
        使用 VL 模型从三级价格详情页截图中识别分时段电价
        """
        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return None

        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        try:
            response = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                timeout=30,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": cls._prompt_price_detail},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},
                            },
                        ],
                    }
                ],
                max_tokens=2000,
                temperature=0.01
            )

            content = response.choices[0].message.content
            logger.info(f"VL Price Detail Response: {content}")

            json_str = cls._extract_json(content)
            data = json.loads(json_str)
            
            # 兼容性处理：如果返回的是对象且包含列表字段，提取列表
            if isinstance(data, dict):
                for key in ["price_list", "prices", "schedule", "data", "items"]:
                    if key in data and isinstance(data[key], list):
                        return data[key]
                
                # 如果字典本身就是一个价格项（包含 start 和 price），将其包装成列表
                if "start" in data and ("price" in data or "total_price" in data):
                    return [data]
            
            # 如果是列表，直接返回
            if isinstance(data, list):
                return data
            
            # 兜底：确保返回的是列表
            return [data] if data else []

        except Exception as e:
            logger.error(f"Error calling VL model for price detail: {e}")
            return None

    @classmethod
    async def get_station_detail_from_image(cls, image_path: str):
        """
        使用 VL 模型从详情页截图中识别地址、名称和电价
        """
        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return None

        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        try:
            response = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                timeout=30,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": cls._prompt_detail},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},
                            },
                        ],
                    }
                ],
                max_tokens=1000,
                temperature=0.01
            )

            content = response.choices[0].message.content
            logger.info(f"VL Detail Response: {content}")

            json_str = cls._extract_json(content)
            detail = json.loads(json_str)
            
            return detail

        except Exception as e:
            logger.error(f"Error calling VL model for detail: {e}")
            return None

    @staticmethod
    async def parse_vl_image(image_path: str, json_metadata, device_info=None):
        """
        基于 _vl.jpg (带绿框) 和 JSON 元数据进行整页识别 (节约 TOKEN)
        """
        if not json_metadata or "cards" not in json_metadata:
            return []

        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return []

        with open(image_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("utf-8")
            img_input = {"url": f"data:image/jpeg;base64,{encoded}"}

        cards_meta = json_metadata["cards"]
        # 按 id 或 Y 坐标排序，确保顺序一致
        cards_meta.sort(key=lambda x: x["rect"][1])

        prompt = (
            "图片中用绿色矩形框标记了若干个充电站卡片区域。\n"
            "请按从上到下的顺序，依次识别每个绿色框内的场站信息，并返回一个JSON数组。\n"
            "数组中元素的顺序必须与图片中绿色框从上到下的顺序严格一致。\n"
            "如果某个框内不是有效的场站卡片（例如是广告或推荐标签），请返回null或空对象，不要跳过顺序。\n"
            "\n"
            "每个JSON对象包含以下字段：\n"
            "1. b_use: 状态标识（1或0）。如果场站名称为灰色或带有“暂停使用”等标签，则为0，否则为1。\n"
            "2. station_name: 场站名称；\n"
            "3. price: 价格（数字，如 1.2500）；\n"
            "4. piles: 充电枪列表 [{type: '快充'|'慢充'|'超充', free: 4, total: 4}]。如果图中只显示数字（如 '34/42'）而没有明确标明类型，请统一识别为 `type: \"慢充\"`；\n"
            "5. parking: 停车费用描述（通常在蓝色'P'图标后，例如 '免费停车三小时'）；\n"
            "6. distance: 距离信息字符串（例如 '2.82km' 或 '90m'）。\n"
            "\n"
            "严格返回纯JSON格式。"
        )

        try:
            resp = await ReadImageKit._client.chat.completions.create(
                model=VL_MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": img_input},
                            {"type": "text", "text": prompt},
                        ],
                    },
                ],
            )

            content = resp.choices[0].message.content or ""
            raw = ReadImageKit._extract_json(content)
            llm_data = json.loads(raw)

            final_stations = []
            if isinstance(llm_data, list):
                # 尝试与 cards_meta 对齐
                count = min(len(llm_data), len(cards_meta))
                for i in range(count):
                    item = llm_data[i]
                    meta = cards_meta[i]

                    if not item or not isinstance(item, dict):
                        continue

                    name = item.get("station_name")
                    if not name or name == "未知":
                        continue

                    # 注入元数据中的点击坐标
                    click_pt = meta.get("click_point", [0, 0])
                    item["uia_center_x"] = click_pt[0]
                    item["uia_center_y"] = click_pt[1]
                    item["rect"] = meta.get("rect")
                    
                    # 转换 bounds 到 0-1000 空间 (用于匹配)
                    img_w = 1080
                    img_h = 2400
                    if device_info:
                        img_w = device_info.get("displayWidth") or device_info.get("width") or 1080
                        img_h = device_info.get("displayHeight") or device_info.get("height") or 2400
                    
                    rect = meta["rect"]
                    item["bounds"] = [
                        int(rect[0] * 1000 / img_w),
                        int(rect[1] * 1000 / img_h),
                        int(rect[2] * 1000 / img_w),
                        int(rect[3] * 1000 / img_h)
                    ]

                    final_stations.append(item)
            
            # 后置过滤
            if final_stations:
                processed_excluded_titles = {str(x).replace(" ", "").strip() for x in STATION_EXCLUDED_TITLES}
                filtered = []
                for item in final_stations:
                    name = item.get("station_name")
                    normalized_name = str(name).replace(" ", "").strip()
                    if normalized_name in processed_excluded_titles:
                        continue
                    if any(kw in normalized_name for kw in STATION_BLACKLIST_KEYWORDS):
                        continue
                    filtered.append(item)
                final_stations = filtered

            return final_stations

        except Exception as e:
            logger.error(f"Error calling VL model for page recognition: {e}")
            return []

    @classmethod
    async def find_time_price_button_coordinate(cls, image_path: str, device_info=None):
        """
        使用 VL 模型定位详情页中的 '分时价格' 按钮
        """
        if device_info is None:
            logger.warning("未提供动态设备信息，使用通用回退配置。")
            device_info = cls._FALLBACK_DEVICE_INFO

        if not os.path.exists(image_path):
            return {}

        with open(image_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("utf-8")
            img_input = {"url": f"data:image/jpeg;base64,{encoded}"}

        prompt = (
            "请找到图片中 '分时价格' 或 '查看分时电价' 按钮区域（通常在详情页，可能是一个带有价格图标或箭头的小型矩形区域）。\n"
            "返回格式示例：\n"
            "{\n"
            '  "bounds": {"x1": 100, "y1": 200, "x2": 300, "y2": 400}, \n'
            '  "bounds_norm": {"left": 0.1, "top": 0.2, "right": 0.3, "bottom": 0.4}\n'
            "}\n"
            "注意：bounds应使用0-1000的归一化坐标空间。\n"
            "如果未找到，返回空JSON {}。"
        )

        try:
            resp = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                timeout=30,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": img_input},
                            {"type": "text", "text": json.dumps(device_info, ensure_ascii=False)},
                            {"type": "text", "text": prompt}
                        ]
                    }
                ],
                max_tokens=200,
                temperature=0.01
            )
            content = resp.choices[0].message.content.strip()
            raw = cls._extract_json(content)
            
            try:
                data = json.loads(raw)
                if isinstance(data, dict) and (data.get("bounds") or data.get("bounds_norm")):
                    data = cls._add_center(data, device_info)
                    # 只返回中心坐标
                    return {
                        "uia_center_x": data.get("uia_center_x"),
                        "uia_center_y": data.get("uia_center_y")
                    }
                return {}
            except Exception as e:
                logger.error(f"Error parsing JSON: {e}")
                return {}
                
        except Exception as e:
            logger.error(f"VL find_time_price_button_coordinate error: {e}")
            return {}