aiData/Apps/TelaiDian/Kit.py

# coding=utf-8
import os
import logging

import hashlib
import numpy as np
import cv2

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

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 detect_cards_cv(image_path, top_ratio=0.15, bottom_ratio=0.1):
    """
    使用计算机图形学 (OpenCV) 检测列表中的场站卡片
    """
    img = read_image(image_path)
    if img is None:
        return []
        
    h, w = img.shape[:2]
    
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blurred, 10, 50)
    
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    cards = []
    min_card_width = int(w * 0.8) 
    min_card_height = 150         
    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)
        if (cw > min_card_width and 
            ch > min_card_height and 
            y > top_limit and 
            y + ch < bottom_limit):
            
            center_y = y + ch // 2
            is_duplicate = False
            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 = 5
                cards.append([
                    max(0, x + padding), 
                    max(0, y + padding), 
                    max(0, x + cw - padding), 
                    max(0, y + ch - padding)
                ])
    
    cards.sort(key=lambda c: c[1])
    return cards

def draw_rectangles(image_path, points, output_path=None):
    """
    使用 OpenCV 在图片上绘制矩形框
    """
    try:
        from Apps.TelaiDian.Config.Setting import DRAW_DEBUG_BOXES
        # 这里暂时写死颜色和粗细，或者从 Setting 读取
        DEBUG_BOX_COLOR = (0, 255, 0)
        DEBUG_BOX_THICKNESS = 3
        
        # 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:
                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:
                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

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}")