aiData/Apps/YeLiTe/Kit.py

# 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
    success, encoded_img = cv2.imencode(".jpg", content)
    if success:
        return hashlib.md5(encoded_img.tobytes()).hexdigest()
    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 detect_cards_cv(image_path, top_ratio=0.15, bottom_ratio=0.1):
    """
    使用计算机图形学 (OpenCV) 检测列表中的场站卡片
    方案：以 nav.jpg 为锚点，向上/下探测白色卡片边界
    """
    img = read_image(image_path)
    if img is None:
        return []
        
    h, w = img.shape[:2]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 1. 模板匹配查找锚点 (nav.jpg)
    template_path = os.path.join(os.path.dirname(__file__), "Template", "nav.jpg")
    anchors = []
    if os.path.exists(template_path):
        template = read_image(template_path)
        if template is not None:
            res = cv2.matchTemplate(img, template, cv2.TM_CCOEFF_NORMED)
            threshold = 0.8
            loc = np.where(res >= threshold)
            
            # 去重记录锚点中心
            t_h, t_w = template.shape[:2]
            matched_points = []
            for pt in zip(*loc[::-1]): # (x, y)
                is_duplicate = False
                for mx, my in matched_points:
                    if abs(mx - pt[0]) < 50 and abs(my - pt[1]) < 50:
                        is_duplicate = True
                        break
                if not is_duplicate:
                    matched_points.append(pt)
                    anchors.append((pt[0] + t_w // 2, pt[1] + t_h // 2))
                    
    logger.info(f"通过模板匹配找到 {len(anchors)} 个锚点")

    # 2. 基于锚点向上/下扩张探测边界
    # 计算每一行的平均亮度，用于探测背景
    row_means = np.mean(gray[:, int(w*0.2):int(w*0.8)], axis=1)
    
    cards = []
    for ax, ay in anchors:
        # 向上探测：直到亮度低于 245 (进入灰色背景)
        y1 = ay
        while y1 > h * top_ratio:
            if row_means[int(y1)] < 245:
                break
            y1 -= 1
        
        # 向下探测：直到亮度低于 245
        y2 = ay
        while y2 < h * (1 - bottom_ratio):
            if row_means[int(y2)] < 245:
                break
            y2 += 1
            
        # 验证卡片高度是否合理
        if y2 - y1 > 150:
            # 给 y 坐标增加一点 padding
            cards.append([int(w * 0.05), int(y1 + 5), int(w * 0.95), int(y2 - 5)])

    # 3. 兜底方案：如果模板匹配没找到，使用之前的投影法
    if not cards:
        logger.warning("模板匹配未找到卡片，尝试投影法兜底...")
        _, binary = cv2.threshold(gray, 250, 255, cv2.THRESH_BINARY)
        top = int(h * top_ratio)
        bottom = int(h * (1 - bottom_ratio))
        roi_binary = binary[:, int(w*0.2):int(w*0.8)]
        row_white_counts = np.sum(roi_binary == 255, axis=1)
        threshold_count = (int(w*0.8) - int(w*0.2)) * 0.9
        is_card_row = row_white_counts > threshold_count
        
        start_y = None
        for y in range(top, bottom):
            if is_card_row[y]:
                if start_y is None: start_y = y
            else:
                if start_y is not None:
                    if y - start_y > 150:
                        cards.append([int(w * 0.05), start_y + 10, int(w * 0.95), y - 10])
                    start_y = None
        if start_y is not None and bottom - start_y > 150:
            cards.append([int(w * 0.05), start_y + 10, int(w * 0.95), bottom - 10])

    # 4. 排序
    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