365 lines
12 KiB
Python
365 lines
12 KiB
Python
# coding=utf-8
|
||
import os
|
||
import time
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
def clean_station_name(name):
|
||
"""
|
||
清理场站名称，去除特殊字符和距离信息
|
||
"""
|
||
if not name:
|
||
return ""
|
||
# 驿来特可能特有的清理逻辑
|
||
import re
|
||
# 移除常见的括号备注
|
||
name = re.sub(r'\(.*?\)', '', name)
|
||
name = re.sub(r'（.*?）', '', name)
|
||
return name.strip()
|
||
|
||
def take_screenshot(d, filename, save_dir=None):
|
||
"""
|
||
获取屏幕截图并保存
|
||
"""
|
||
if not save_dir:
|
||
from Config.Config import TEMP_IMAGE_DIR
|
||
save_dir = TEMP_IMAGE_DIR
|
||
|
||
if not os.path.exists(save_dir):
|
||
os.makedirs(save_dir)
|
||
|
||
# 确保文件名有后缀
|
||
if not filename.endswith(".jpg") and not filename.endswith(".png"):
|
||
filename = f"{filename}.jpg"
|
||
|
||
full_path = os.path.join(save_dir, filename)
|
||
d.screenshot(full_path)
|
||
return full_path
|
||
|
||
import hashlib
|
||
import numpy as np
|
||
import cv2
|
||
|
||
def get_file_md5(file_path):
|
||
"""计算文件的 MD5 值"""
|
||
if not os.path.exists(file_path):
|
||
return None
|
||
hash_md5 = hashlib.md5()
|
||
with open(file_path, "rb") as f:
|
||
for chunk in iter(lambda: f.read(4096), b""):
|
||
hash_md5.update(chunk)
|
||
return hash_md5.hexdigest()
|
||
|
||
def get_image_content_md5(file_path, top_ratio=0.1, bottom_ratio=0.1):
|
||
"""
|
||
计算图片核心内容的 MD5 值（排除状态栏和导航栏）
|
||
"""
|
||
img = read_image(file_path)
|
||
if img is None:
|
||
return None
|
||
|
||
h, w = img.shape[:2]
|
||
top = int(h * top_ratio)
|
||
bottom = int(h * (1 - bottom_ratio))
|
||
|
||
# 裁剪中间部分
|
||
content = img[top:bottom, :]
|
||
|
||
# 直接使用原始像素数组计算 MD5，避免压缩损失
|
||
return hashlib.md5(content.tobytes()).hexdigest()
|
||
|
||
def read_image(path):
|
||
"""读取图片，支持中文路径"""
|
||
if not path or not os.path.exists(path):
|
||
return None
|
||
try:
|
||
data = np.fromfile(path, dtype=np.uint8)
|
||
if data.size == 0:
|
||
return None
|
||
img = cv2.imdecode(data, -1)
|
||
return img
|
||
except Exception as e:
|
||
logger.error(f"Error reading image {path}: {e}")
|
||
return None
|
||
|
||
def save_image(path, img):
|
||
"""保存图片，支持中文路径"""
|
||
try:
|
||
ext = os.path.splitext(path)[1]
|
||
if not ext:
|
||
ext = ".jpg"
|
||
cv2.imencode(ext, img)[1].tofile(path)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Error saving image {path}: {e}")
|
||
return False
|
||
|
||
def find_template(image_path, template_name, threshold=0.8):
|
||
"""
|
||
在图片中查找模板位置
|
||
"""
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return None
|
||
|
||
template_path = os.path.join(os.path.dirname(__file__), "Template", template_name)
|
||
if not os.path.exists(template_path):
|
||
logger.warning(f"Template not found: {template_path}")
|
||
return None
|
||
|
||
template = read_image(template_path)
|
||
if template is None:
|
||
return None
|
||
|
||
res = cv2.matchTemplate(img, template, cv2.TM_CCOEFF_NORMED)
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
|
||
|
||
if max_val >= threshold:
|
||
t_h, t_w = template.shape[:2]
|
||
center_x = max_loc[0] + t_w // 2
|
||
center_y = max_loc[1] + t_h // 2
|
||
return (center_x, center_y)
|
||
return None
|
||
|
||
def detect_cards_cv(image_path, top_ratio=0.15, bottom_ratio=0.1):
|
||
"""
|
||
使用计算机图形学 (OpenCV) 检测列表中的场站卡片
|
||
方案：使用 Canny 边缘检测 + 轮廓查找，提取圆角矩形卡片
|
||
"""
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return []
|
||
|
||
h, w = img.shape[:2]
|
||
|
||
# 1. 预处理：转灰度 -> 高斯模糊去噪
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
||
|
||
# 2. Canny 边缘检测
|
||
# 阈值设置：降低阈值以适应低对比度边缘 (如浅灰背景上的白卡片)
|
||
# 经过测试，(10, 50) 对于浅灰背景上的白卡片效果最佳
|
||
edges = cv2.Canny(blurred, 10, 50)
|
||
|
||
# 3. 查找轮廓
|
||
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
cards = []
|
||
min_card_width = int(w * 0.8) # 卡片宽度通常占屏幕 80% 以上
|
||
min_card_height = 150 # 卡片高度至少 150px
|
||
top_limit = int(h * top_ratio)
|
||
bottom_limit = int(h * (1 - bottom_ratio))
|
||
|
||
valid_contours = []
|
||
|
||
for cnt in contours:
|
||
x, y, cw, ch = cv2.boundingRect(cnt)
|
||
|
||
# 筛选条件：
|
||
# 1. 宽度足够大 (排除小图标、文字)
|
||
# 2. 高度适中
|
||
# 3. 位置在安全区域内 (排除顶部导航和底部 TabBar)
|
||
if (cw > min_card_width and
|
||
ch > min_card_height and
|
||
y > top_limit and
|
||
y + ch < bottom_limit):
|
||
|
||
# 简单去重：检查是否与已有的卡片高度重叠
|
||
is_duplicate = False
|
||
center_y = y + ch // 2
|
||
for vx, vy, vcw, vch in valid_contours:
|
||
v_center_y = vy + vch // 2
|
||
if abs(center_y - v_center_y) < 50: # 中心点接近，视为同一个
|
||
is_duplicate = True
|
||
break
|
||
|
||
if not is_duplicate:
|
||
valid_contours.append((x, y, cw, ch))
|
||
# 增加一点 padding，收缩一下边界，避免点到边缘
|
||
padding = 5
|
||
cards.append([
|
||
max(0, x + padding),
|
||
max(0, y + padding),
|
||
min(w, x + cw - padding),
|
||
min(h, y + ch - padding)
|
||
])
|
||
|
||
# 按 y 坐标排序
|
||
cards.sort(key=lambda c: c[1])
|
||
|
||
# 如果没找到，尝试稍微放宽条件兜底 (比如宽度 > 60%)
|
||
if not cards:
|
||
logger.warning("严格轮廓检测未找到卡片，尝试放宽条件...")
|
||
for cnt in contours:
|
||
x, y, cw, ch = cv2.boundingRect(cnt)
|
||
if (cw > int(w * 0.6) and ch > 100 and y > top_limit and y + ch < bottom_limit):
|
||
# 简单去重
|
||
center_y = y + ch // 2
|
||
is_duplicate = False
|
||
for existing in cards:
|
||
ec_y = (existing[1] + existing[3]) // 2
|
||
if abs(center_y - ec_y) < 50:
|
||
is_duplicate = True
|
||
break
|
||
if not is_duplicate:
|
||
cards.append([x, y, x+cw, y+ch])
|
||
cards.sort(key=lambda c: c[1])
|
||
|
||
logger.info(f"最终检测到 {len(cards)} 个场站卡片区域")
|
||
return cards
|
||
|
||
def clear_temp_dir(save_dir=None):
|
||
"""清空临时目录中的所有文件"""
|
||
if save_dir is None:
|
||
from Config.Config import TEMP_IMAGE_DIR
|
||
save_dir = TEMP_IMAGE_DIR
|
||
|
||
if not os.path.exists(save_dir):
|
||
return
|
||
|
||
for f in os.listdir(save_dir):
|
||
file_path = os.path.join(save_dir, f)
|
||
try:
|
||
if os.path.isfile(file_path):
|
||
os.remove(file_path)
|
||
except Exception as e:
|
||
logger.error(f"Error deleting file {file_path}: {e}")
|
||
|
||
def get_image_md5(image_path):
|
||
"""
|
||
计算图片的 MD5 值
|
||
"""
|
||
import hashlib
|
||
if not os.path.exists(image_path):
|
||
return ""
|
||
with open(image_path, 'rb') as f:
|
||
data = f.read()
|
||
return hashlib.md5(data).hexdigest()
|
||
|
||
def draw_rectangles(image_path, points, output_path=None):
|
||
"""
|
||
使用 OpenCV 在图片上绘制矩形框 (计算机图形学绘制)
|
||
:param image_path: 原始图片路径
|
||
:param points: 坐标点列表，格式如 [[x1, y1, x2, y2], ...] 或 [(x, y), ...]
|
||
:param output_path: 输出图片路径，如果不指定则覆盖原图
|
||
"""
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
from Apps.YeLiTe.Config.Setting import DRAW_DEBUG_BOXES, DEBUG_BOX_COLOR, DEBUG_BOX_THICKNESS
|
||
|
||
if not DRAW_DEBUG_BOXES:
|
||
return image_path
|
||
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return image_path
|
||
|
||
for p in points:
|
||
if len(p) == 4: # 矩形 [x1, y1, x2, y2]
|
||
cv2.rectangle(img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), DEBUG_BOX_COLOR, DEBUG_BOX_THICKNESS)
|
||
elif len(p) == 2: # 点 (x, y)，画一个小圆圈或正方形表示点击位
|
||
center = (int(p[0]), int(p[1]))
|
||
cv2.circle(img, center, 10, DEBUG_BOX_COLOR, -1)
|
||
|
||
save_path = output_path if output_path else image_path
|
||
save_image(save_path, img)
|
||
return save_path
|
||
except Exception as e:
|
||
logger.error(f"绘制矩形框失败: {e}")
|
||
return image_path
|
||
|
||
from Util.EasyOcrKit import get_easyocr_reader
|
||
|
||
# 预加载 EasyOCR Reader (单例模式)
|
||
def get_ocr_reader():
|
||
return get_easyocr_reader(gpu=True)
|
||
|
||
def detect_price_info_container_cv(image_path):
|
||
"""
|
||
使用 OCR 精准定位详情页中的价格入口文本（“阶段性电价”）。
|
||
返回: [x1, y1, x2, y2] 绝对像素坐标，如果未找到则返回 None
|
||
"""
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return None
|
||
h, w = img.shape[:2]
|
||
|
||
keywords = ['阶段性电价']
|
||
|
||
try:
|
||
reader = get_ocr_reader()
|
||
results = reader.read_text(img)
|
||
|
||
for (quad, text, prob) in results:
|
||
if any(kw in text for kw in keywords) and prob >= 0.5:
|
||
pts = np.array(quad).astype(int)
|
||
x_min = np.min(pts[:, 0])
|
||
y_min = np.min(pts[:, 1])
|
||
x_max = np.max(pts[:, 0])
|
||
y_max = np.max(pts[:, 1])
|
||
|
||
center_x = (x_min + x_max) // 2
|
||
center_y = (y_min + y_max) // 2
|
||
logger.info(f"[OCR识别] 找到文本: '{text}', 置信度: {prob:.4f}, 中心坐标: ({center_x}, {center_y})")
|
||
return (center_x, center_y)
|
||
except Exception as e:
|
||
logger.error(f"OCR 识别发生异常: {e}")
|
||
|
||
return None
|
||
|
||
def setup_logger(name, log_file=None, clear_old_log=False):
|
||
"""
|
||
配置日志，支持同时输出到控制台和文件。
|
||
使用供应商代号作为父级 Logger，所有子 Logger 继承其 Handler，
|
||
并通过 propagate=False 避免与根 Logger 重复。
|
||
:param name: Logger 名称
|
||
:param log_file: 指定日志文件路径，如果不指定则使用默认路径
|
||
:param clear_old_log: 是否在启动时清空旧日志文件
|
||
"""
|
||
# 1. 获取供应商代号 (如 YeLiTe)
|
||
supplier_code = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
# 2. 获取父级 Logger 并配置
|
||
parent_logger = logging.getLogger(supplier_code)
|
||
parent_logger.setLevel(logging.INFO)
|
||
parent_logger.propagate = False # 禁止向上传递给 root logger，防止重复
|
||
|
||
if log_file is None:
|
||
# 获取项目根目录 (aiData)
|
||
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
log_dir = os.path.join(root_dir, "Logs")
|
||
if not os.path.exists(log_dir):
|
||
os.makedirs(log_dir)
|
||
log_file = os.path.join(log_dir, f"{supplier_code}.log")
|
||
|
||
# 如果需要清空旧日志且文件存在
|
||
if clear_old_log and os.path.exists(log_file):
|
||
try:
|
||
# 关闭现有的 handler 以便删除文件
|
||
for handler in parent_logger.handlers[:]:
|
||
handler.close()
|
||
parent_logger.removeHandler(handler)
|
||
os.remove(log_file)
|
||
except Exception as e:
|
||
print(f"无法清空旧日志文件 {log_file}: {e}")
|
||
|
||
if not parent_logger.handlers:
|
||
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
|
||
# 控制台 Handler
|
||
ch = logging.StreamHandler()
|
||
ch.setFormatter(formatter)
|
||
parent_logger.addHandler(ch)
|
||
|
||
# 文件 Handler
|
||
fh = logging.FileHandler(log_file, encoding='utf-8')
|
||
fh.setFormatter(formatter)
|
||
parent_logger.addHandler(fh)
|
||
|
||
# 3. 返回子 Logger
|
||
if name == supplier_code:
|
||
return parent_logger
|
||
return logging.getLogger(f"{supplier_code}.{name}")
|