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
from openai import OpenAI, BadRequestError
from Config.Config import (
    ALY_LLM_API_KEY, VL_MODEL_NAME, VL_MODEL_NAME_AD
)
from Apps.AiTeJiYiChong.Config.Setting import (
    SAFE_EXCLUDE_RATIO, FALLBACK_WIDTH, FALLBACK_HEIGHT,
    BOTTOM_SAFE_EXCLUDE_RATIO
)

logger = logging.getLogger(__name__)


class ReadImageKit:
    _client = OpenAI(
        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，包含'快'或'慢'的类型、总枪数和空闲枪数，例如 '快 4/4')。\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 asyncio.to_thread(
                cls._client.chat.completions.create,
                model=VL_MODEL_NAME,
                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"]:
                    if key in data and isinstance(data[key], list):
                        return data[key]
                # 如果本身就是个包含列表的字典，也可能需要根据实际情况调整
                # 这里假设 prompt 引导它返回一个包含列表的对象或直接是列表
                if "items" in data: return data["items"]
            
            return data

        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 asyncio.to_thread(
                cls._client.chat.completions.create,
                model=VL_MODEL_NAME,
                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

    @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 asyncio.to_thread(
                cls._client.chat.completions.create,
                model=VL_MODEL_NAME,
                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 []