133 lines
5.3 KiB
Python
133 lines
5.3 KiB
Python
# coding=utf-8
|
|
import logging
|
|
import os
|
|
import sys
|
|
import json
|
|
|
|
# Ensure sys path includes root for imports if not already
|
|
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
if project_root not in sys.path:
|
|
sys.path.append(project_root)
|
|
|
|
from Util.VLMKit import VLMKit
|
|
from Apps.TelaiDian.Kit import draw_rectangles, detect_cards_cv, setup_logger
|
|
from Apps.TelaiDian.Config.Setting import SAFE_EXCLUDE_RATIO, BOTTOM_SAFE_EXCLUDE_RATIO
|
|
|
|
# 初始化日志
|
|
logger = setup_logger("ReadImageKit")
|
|
|
|
class ReadImageKit:
|
|
def __init__(self):
|
|
self.vlm = VLMKit()
|
|
|
|
async def analyze_detail_price(self, image_path):
|
|
"""
|
|
分析详情页截图，提取电价信息
|
|
"""
|
|
prompt = """
|
|
分析这张充电站详情页截图，提取**电价时段表**。
|
|
请仔细寻找包含“时段”、“电价”、“服务费”或“总价”的表格或列表。
|
|
|
|
请提取每个时段的：
|
|
1. 开始时间 (HH:MM)
|
|
2. 结束时间 (HH:MM)
|
|
3. 总电价 (元/度，包含电费和服务费)
|
|
|
|
输出格式为 JSON 数组：
|
|
[
|
|
{
|
|
"start": "00:00",
|
|
"end": "08:00",
|
|
"price": 1.23
|
|
},
|
|
...
|
|
]
|
|
如果无法识别任何价格信息，请返回空数组 []。
|
|
"""
|
|
try:
|
|
res_text = await self.vlm.analyze_image(image_path, prompt)
|
|
logger.info(f"VLM Price Analysis Result for {os.path.basename(image_path)}: {res_text[:200]}...")
|
|
|
|
json_str = self.vlm.extract_json(res_text)
|
|
prices = json.loads(json_str)
|
|
|
|
normalized_prices = []
|
|
if isinstance(prices, list):
|
|
for p in prices:
|
|
new_p = p.copy()
|
|
if 'time_range' in p and ('start' not in p or 'end' not in p):
|
|
tr = p['time_range'].replace('~', '-').replace(' ', '')
|
|
parts = tr.split('-')
|
|
if len(parts) >= 2:
|
|
new_p['start'] = parts[0]
|
|
new_p['end'] = parts[1]
|
|
|
|
if 'price' not in p:
|
|
if 'total_price' in p:
|
|
new_p['price'] = p['total_price']
|
|
elif 'elec_price' in p and 'service_price' in p:
|
|
try:
|
|
new_p['price'] = float(p['elec_price']) + float(p['service_price'])
|
|
except:
|
|
pass
|
|
normalized_prices.append(new_p)
|
|
return normalized_prices
|
|
return []
|
|
except Exception as e:
|
|
logger.error(f"分析电价详情失败: {e}")
|
|
return []
|
|
|
|
async def analyze_station_list(self, image_path):
|
|
"""
|
|
分析场站列表页图片，提取场站位置和基本信息
|
|
"""
|
|
cv_bboxes = detect_cards_cv(image_path, top_ratio=SAFE_EXCLUDE_RATIO, bottom_ratio=BOTTOM_SAFE_EXCLUDE_RATIO)
|
|
|
|
if cv_bboxes:
|
|
draw_rectangles(image_path, cv_bboxes)
|
|
prompt = f"""
|
|
图片中已经用绿色矩形框标记了 {len(cv_bboxes)} 个可能的充电站卡片。
|
|
请按从上到下的顺序，识别每个绿色框内的场站信息。
|
|
|
|
输出格式为 JSON 数组，长度必须为 {len(cv_bboxes)}。
|
|
每个对象包含：
|
|
- "name": 场站名称
|
|
- "address": 场站地址
|
|
- "is_valid": true/false (是否为真实的场站卡片)
|
|
"""
|
|
else:
|
|
prompt = """
|
|
分析这张充电站列表截图，提取所有充电站卡片信息。
|
|
忽略顶部的筛选栏，仅提取下方重复出现的场站卡片。
|
|
输出格式为 JSON 数组，每个对象包含：
|
|
- "name": 场站名称
|
|
- "address": 场站地址
|
|
- "point": 场站卡片的中心点击坐标 [x, y]
|
|
- "bbox": 场站卡片的边界框 [x1, y1, x2, y2]
|
|
"""
|
|
|
|
try:
|
|
res_text = await self.vlm.analyze_image(image_path, prompt)
|
|
json_str = self.vlm.extract_json(res_text)
|
|
vlm_results = json.loads(json_str)
|
|
|
|
final_stations = []
|
|
if cv_bboxes and isinstance(vlm_results, list):
|
|
for i, res in enumerate(vlm_results):
|
|
if i < len(cv_bboxes):
|
|
bbox = cv_bboxes[i]
|
|
if res and (res.get("is_valid") is True or (res.get("name") and res.get("is_valid") is not False)):
|
|
final_stations.append({
|
|
"name": res.get("name"),
|
|
"address": res.get("address"),
|
|
"point": [(bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2],
|
|
"bbox": bbox
|
|
})
|
|
elif not cv_bboxes:
|
|
final_stations = vlm_results if isinstance(vlm_results, list) else []
|
|
|
|
return final_stations
|
|
except Exception as e:
|
|
logger.error(f"分析列表页失败: {e}")
|
|
return []
|