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
)

logger = logging.getLogger(__name__)

# 场站过滤黑名单
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"
    }

    _prompt = (
        "仅输出JSON数组（不含任何说明文字），按从左到右、从上到下的顺序识别图片中由【绿色方框】标识的充电站区域。识别规则如下：\n"
        "1. 必须是图中用绿色实线方框圈出的区域。\n"
        "2. 每一个卡片区域必须同时具备以下所有要素，否则严禁识别：\n"
        "   - 场站名称 (station_name)；\n"
        "   - 距离信息 (distance, 例如 '2.82km' 或 '90m')，通常位于卡片右侧蓝色胶囊区域内；\n"
        "   - 金额/电费 (price，例如 '1.2500')，通常以红色字体显示；\n"
        "   - **充电枪信息 (piles)**: 必须包含类型('快充'、'慢充'、'超充')、总枪数和空闲枪数。识别规则如下：\n"
        "     - 如果明确标有'快'、'慢'或'超'，则对应为'快充'、'慢充'、'超充'。\n"
        "     - 如果图中只显示数字（如 '34/42'）而没有明确标明类型，请统一识别为 `type: \"慢充\"`。\n"
        "     - 必须以 JSON 对象列表形式输出，例如：`[{\"type\": \"快充\", \"free\": 4, \"total\": 4}, {\"type\": \"慢充\", \"free\": 1, \"total\": 2}]`。\n"
        "3. 如果绿色方框内缺少上述任何一项要素，说明它不是真正的场站卡片，请直接跳过。\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}]；\n"
        "5. parking: 停车费用描述（通常在蓝色'P'图标后，例如 '免费停车三小时'）；\n"
        "6. distance: 距离信息字符串（例如 '2.82km' 或 '90m'）；\n"
        "7. bounds: {x1,y1,x2,y2} 区域像素坐标（0-1000）；\n"
        "8. bounds_norm: {left,top,right,bottom} 归一化坐标(0-1)；\n"
        "9. station_name_bounds: 场站名称文字区域坐标 {x1,y1,x2,y2}（0-1000）；\n"
        "10. station_name_bounds_norm: 场站名称文字归一化坐标(0-1)。\n"
        "\n"
        "重要约束：\n"
        "A. 严禁识别未被绿色方框圈出的区域。如顶部的“长春市”选择、顶部的搜索框、以及中间的“推荐站点”等标签。\n"
        "B. 真正的场站卡片在绿色方框内包含：场站名称、金额（红色文字）、距离（右侧蓝色背景内）、充电枪状态（绿色或蓝色徽章，格式为 闲x/x 或 x/x）。\n"
        "C. 严禁将顶部的功能图标（如我的订单、收藏站点等）误认为场站卡片。\n"
        "\n"
        "严格返回纯JSON格式。"
    )

    @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对象。"
    )

    @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 recognize_card_with_vl(image_path: str):
        """
        对卡片切片进行 VL 识别
        """
        if not os.path.exists(image_path):
            return None

        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"
            "1. **场站名称 (station_name)**: 场站名称通常位于卡片的第一行（左上角），字体最大且加粗。严禁提取'距离优先'、'空闲较多'、'价格最低'等筛选标签。严禁将卡片中下部的描述信息（如'导航到...'、'停车说明...'、'距您...'等）识别为名称。名称通常是一个具体的地点或品牌名。\n"
            "2. **价格 (price)**: 提取数字部分，通常以'元/度'结尾，字体颜色通常为红色或加粗。\n"
            "3. **桩数 (piles)**: 识别所有类型的充电桩。请提取'快'、'慢'或'超'字样及其对应的'空闲/总数'（如 '快 4/4'）。\n"
            "   - 如果明确标有'快'、'慢'或'超'，则对应为'快充'、'慢充'、'超充'。\n"
            "   - 如果图中只显示数字（如 '34/42'）而没有明确标明类型，请统一识别为 `type: \"慢充\"`。\n"
            "   - 必须以 JSON 对象列表形式输出，例如：`[{\"type\": \"快充\", \"free\": 4, \"total\": 4}]`。\n"
            "4. **距离 (distance)**: 提取如 '2.5km' 格式，通常在蓝色胶囊背景内。\n"
            "5. **停车 (parking)**: 简要描述停车收费信息，通常在蓝色'P'图标后。\n\n"
            "请以 JSON 格式输出：\n"
            "{\n"
            "  \"station_name\": \"场站名称\",\n"
            "  \"price\": \"价格\",\n"
            "  \"piles\": [\n"
            "    {\"type\": \"快\", \"free\": 4, \"total\": 4}\n"
            "  ],\n"
            "  \"distance\": \"距离\",\n"
            "  \"parking\": \"停车信息\"\n"
            "}\n"
            "只输出 JSON，不要有任何解释文字。"
        )

        try:
            resp = await ReadImageKit._client.chat.completions.create(
                model=VL_MODEL_NAME,
                timeout=30,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {"type": "image_url", "image_url": img_input}
                        ]
                    }
                ],
                max_tokens=300,
                temperature=0.01
            )
            content = resp.choices[0].message.content.strip()
            raw = ReadImageKit._extract_json(content)
            return json.loads(raw)
        except Exception as e:
            logger.error(f"VL recognize_card_with_vl error: {e}")
            return None

    @staticmethod
    async def get_stations_hybrid(image_path, device_info=None):
        """
        混合识别模式：图形学切片 + VL 模型识别
        """
        if device_info is None:
            device_info = ReadImageKit._FALLBACK_DEVICE_INFO

        # 1. 使用 Kit 中的图形学算法识别卡片区域
        json_data = Kit.crop_cards_from_image(image_path, save_debug=True)
        if not json_data or not json_data.get("cards"):
            logger.warning("No cards detected by graphical slicing.")
            return []

        # 2. 对每个卡片区域进行 VL 识别
        tasks = []
        card_metas = []
        
        # 为了使用 VL，我们需要将每个卡片切出来保存到临时文件
        temp_dir = os.path.join(os.path.dirname(image_path), "Patches")
        os.makedirs(temp_dir, exist_ok=True)
        
        img = Image.open(image_path).convert("RGB")
        
        for card in json_data["cards"]:
            rect = card["rect"] # [x1, y1, x2, y2]
            patch_path = os.path.join(temp_dir, f"patch_{card['id']}_{uuid.uuid4().hex[:8]}.jpg")
            
            # 裁剪并保存
            patch = img.crop((rect[0], rect[1], rect[2], rect[3]))
            patch.save(patch_path)
            
            card_metas.append({
                "meta": card,
                "patch_path": patch_path
            })
            tasks.append(ReadImageKit.recognize_card_with_vl(patch_path))

        # 并行执行 VL 识别
        results = await asyncio.gather(*tasks)
        
        final_stations = []
        for i, res in enumerate(results):
            meta = card_metas[i]["meta"]
            patch_path = card_metas[i]["patch_path"]
            
            if res and isinstance(res, dict) and res.get("station_name"):
                # 注入点击坐标和原始区域信息
                res["uia_center_x"] = meta["click_point"][0]
                res["uia_center_y"] = meta["click_point"][1]
                res["rect"] = meta["rect"]
                
                # 转换 bounds 到 0-1000 空间
                w, h = img.size
                rect = meta["rect"]
                res["bounds"] = [
                    int(rect[0] * 1000 / w),
                    int(rect[1] * 1000 / h),
                    int(rect[2] * 1000 / w),
                    int(rect[3] * 1000 / h)
                ]
                
                final_stations.append(res)
                logger.info(f"卡片 {meta['id']} 识别成功: {res['station_name']}")
            else:
                logger.warning(f"卡片 {meta['id']} 识别失败或无名称")
            
            # 清理临时切片文件
            try:
                if os.path.exists(patch_path):
                    os.remove(patch_path)
            except:
                pass

        # 3. 后置过滤
        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

    @classmethod
    async def get_stations_from_image(cls, image_path: str, device_info=None):
        """
        使用 Qwen-VL 模型从截图中识别充电站列表
        """
        if device_info is None:
            device_info = cls._FALLBACK_DEVICE_INFO

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

        # 将图片转换为 Base64
        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},
                            {
                                "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 Model Response: {content}")

            json_str = cls._extract_json(content)
            stations = json.loads(json_str)

            # 后处理：如果 bounds 是归一化的，则转换为像素坐标（如果需要）
            # 或者如果 bounds 是 0-1000 的，则保持原样或按需转换
            return stations

        except Exception as e:
            logger.error(f"Error calling VL model: {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 {}