Files
aiData/Apps/YeLiTe/Kit.py

308 lines
9.8 KiB
Python
Raw Normal View History

2026-01-13 20:36:20 +08:00
# 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()
2026-01-13 20:46:18 +08:00
def take_screenshot(d, filename, save_dir=None):
2026-01-13 20:36:20 +08:00
"""
获取屏幕截图并保存
"""
2026-01-13 20:46:18 +08:00
if not save_dir:
2026-01-13 20:36:20 +08:00
from Config.Config import TEMP_IMAGE_DIR
2026-01-13 20:46:18 +08:00
save_dir = TEMP_IMAGE_DIR
if not os.path.exists(save_dir):
os.makedirs(save_dir)
2026-01-13 20:36:20 +08:00
2026-01-13 20:46:18 +08:00
# 确保文件名有后缀
if not filename.endswith(".jpg") and not filename.endswith(".png"):
filename = f"{filename}.jpg"
2026-01-13 20:36:20 +08:00
2026-01-13 20:46:18 +08:00
full_path = os.path.join(save_dir, filename)
2026-01-13 20:36:20 +08:00
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, :]
2026-01-13 21:37:14 +08:00
# 直接使用原始像素数组计算 MD5,避免压缩损失
2026-01-13 20:36:20 +08:00
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
2026-01-13 21:37:14 +08:00
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
2026-01-13 20:54:49 +08:00
def detect_cards_cv(image_path, top_ratio=0.15, bottom_ratio=0.1):
"""
使用计算机图形学 (OpenCV) 检测列表中的场站卡片
2026-01-13 21:37:14 +08:00
方案:使用 Canny 边缘检测 + 轮廓查找,提取圆角矩形卡片
2026-01-13 20:54:49 +08:00
"""
img = read_image(image_path)
if img is None:
return []
h, w = img.shape[:2]
2026-01-13 21:37:14 +08:00
# 1. 预处理:转灰度 -> 高斯模糊去噪
2026-01-13 20:54:49 +08:00
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
2026-01-13 21:37:14 +08:00
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
2026-01-13 20:54:49 +08:00
2026-01-13 21:37:14 +08:00
# 2. Canny 边缘检测
# 阈值设置:降低阈值以适应低对比度边缘 (如浅灰背景上的白卡片)
# 经过测试,(10, 50) 对于浅灰背景上的白卡片效果最佳
edges = cv2.Canny(blurred, 10, 50)
# 3. 查找轮廓
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
2026-01-13 20:54:49 +08:00
cards = []
2026-01-13 21:37:14 +08:00
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)
2026-01-13 20:54:49 +08:00
2026-01-13 21:37:14 +08:00
# 筛选条件:
# 1. 宽度足够大 (排除小图标、文字)
# 2. 高度适中
# 3. 位置在安全区域内 (排除顶部导航和底部 TabBar)
if (cw > min_card_width and
ch > min_card_height and
y > top_limit and
y + ch < bottom_limit):
2026-01-13 20:54:49 +08:00
2026-01-13 21:37:14 +08:00
# 简单去重:检查是否与已有的卡片高度重叠
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%)
2026-01-13 20:54:49 +08:00
if not cards:
2026-01-13 21:37:14 +08:00
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])
2026-01-13 20:54:49 +08:00
logger.info(f"最终检测到 {len(cards)} 个场站卡片区域")
return cards
2026-01-13 20:36:20 +08:00
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
2026-01-13 20:46:18 +08:00
from Apps.YeLiTe.Config.Setting import DRAW_DEBUG_BOXES, DEBUG_BOX_COLOR, DEBUG_BOX_THICKNESS
2026-01-13 20:36:20 +08:00
if not DRAW_DEBUG_BOXES:
return image_path
2026-01-13 20:54:49 +08:00
img = read_image(image_path)
2026-01-13 20:36:20 +08:00
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
2026-01-13 20:54:49 +08:00
save_image(save_path, img)
2026-01-13 20:36:20 +08:00
return save_path
except Exception as e:
logger.error(f"绘制矩形框失败: {e}")
return image_path
2026-01-14 07:28:06 +08:00
def setup_logger(name, log_file=None):
"""
配置日志,支持同时输出到控制台和文件
2026-01-14 07:29:32 +08:00
统一存放在根目录的 Logs 文件夹下
2026-01-14 07:28:06 +08:00
"""
if log_file is None:
2026-01-14 07:29:32 +08:00
# 获取项目根目录 (aiData)
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
log_dir = os.path.join(root_dir, "Logs")
2026-01-14 07:28:06 +08:00
if not os.path.exists(log_dir):
os.makedirs(log_dir)
2026-01-14 07:29:32 +08:00
# 统一使用供应商代号作为文件名
supplier_code = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
log_file = os.path.join(log_dir, f"{supplier_code}.log")
2026-01-14 07:28:06 +08:00
logger = logging.getLogger(name)
logger.setLevel(logging.INFO)
# 如果已经有 handler,不再添加
if logger.handlers:
return logger
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
# 控制台 Handler
ch = logging.StreamHandler()
ch.setFormatter(formatter)
logger.addHandler(ch)
# 文件 Handler
fh = logging.FileHandler(log_file, encoding='utf-8')
fh.setFormatter(formatter)
logger.addHandler(fh)
return logger