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_list_price_blocks_cv(image_path, top_ratio=0.40, bottom_ratio=0.12):
    img = read_image(image_path)
    if img is None:
        return []
    h, w = img.shape[:2]
    top_limit = int(h * top_ratio)
    bottom_limit = int(h * (1 - bottom_ratio))
    roi_bgr = img[top_limit:bottom_limit, :, :]
    if roi_bgr.size == 0:
        return []
    b, g, r = cv2.split(roi_bgr)
    mask_bgr = (r > 190) & (g > 80) & (r - g > 25) & (r - b > 25)
    mask_bgr = (mask_bgr.astype(np.uint8)) * 255
    if cv2.countNonZero(mask_bgr) < 50:
        hsv = cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV)
        lower_red1 = np.array([0, 100, 120])
        upper_red1 = np.array([20, 255, 255])
        lower_red2 = np.array([160, 100, 120])
        upper_red2 = np.array([180, 255, 255])
        mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
        mask = cv2.bitwise_or(mask1, mask2)
    else:
        mask = mask_bgr
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 10))
    closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    blocks = []
    for cnt in contours:
        x, y, cw, ch = cv2.boundingRect(cnt)
        y_global = y + top_limit
        if ch < 30 or ch > 140:
            continue
        if cw < 80 or cw > int(w * 0.8):
            continue
        center_y = y_global + ch // 2
        if center_y < int(h * 0.6) or center_y > bottom_limit:
            continue
        blocks.append([x, y_global, x + cw, y_global + ch])
    blocks.sort(key=lambda b: b[1])
    return blocks

def detect_price_click_point_cv(image_path):
    """
    使用 HSV 颜色过滤定位详情页的橘红色价格区域，返回最左侧区域的中心点击点
    """
    img = read_image(image_path)
    if img is None:
        return None
    
    h, w = img.shape[:2]
    # 1. 转换为 HSV 空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # 2. 橘红色的 HSV 范围 (适配特来电价格颜色)
    lower_orange = np.array([0, 150, 150])
    upper_orange = np.array([20, 255, 255])
    mask = cv2.inRange(hsv, lower_orange, upper_orange)
    
    # 3. 对掩码进行膨胀，连接数字
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 20))
    dilated = cv2.dilate(mask, kernel)
    
    # 4. 寻找轮廓
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    detected_areas = []
    for cnt in contours:
        x, y, cw, ch = cv2.boundingRect(cnt)
        # 1. 过滤掉宽度过大（可能是横幅广告）或过小（可能是杂点）的区域
        # 2. 价格区域 P0 通常在屏幕的中部，且宽度约为屏幕的一半
        if 200 < y < h * 0.8 and 100 < cw < w * 0.6 and ch > 30:
            detected_areas.append([x, y, x + cw, y + ch])

    if not detected_areas:
        # 备选：如果 HSV 失败，尝试通过轮廓大小寻找
        # 1.1556 这种大数字通常会有很明显的轮廓
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 10))
        closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            x, y, cw, ch = cv2.boundingRect(cnt)
            if 200 < y < h * 0.6 and 150 < cw < 300 and 50 < ch < 150:
                detected_areas.append([x, y, x + cw, y + ch])

    if not detected_areas:
        return None

    # 5. 按 X 轴排序，取最左边的区域 (即用户确认的 P0)
    # 但要排除掉可能在最左侧的导航栏返回按钮等小元素，所以前面加了宽度限制
    detected_areas.sort(key=lambda b: b[0])
    target = detected_areas[0]
    
    center_x = (target[0] + target[2]) // 2
    center_y = (target[1] + target[3]) // 2
    
    return [center_x, center_y]

def detect_cards_cv(image_path, top_ratio=0.40, bottom_ratio=0.12):
    """
    使用计算机图形学 (OpenCV) 检测列表中的场站卡片。
    """
    from Apps.TeLaiDian.Config.Setting import MIN_CARD_HEIGHT
    
    img = read_image(image_path)
    if img is None:
        return []
        
    h, w = img.shape[:2]
        
    price_blocks = detect_list_price_blocks_cv(image_path, top_ratio=top_ratio, bottom_ratio=bottom_ratio)
    cards = []
    if price_blocks:
        from Apps.TeLaiDian.Config.Setting import SAFE_EXCLUDE_RATIO, BOTTOM_SAFE_EXCLUDE_RATIO
        top_limit = int(h * SAFE_EXCLUDE_RATIO)
        bottom_limit = int(h * (1 - BOTTOM_SAFE_EXCLUDE_RATIO))
        for b in price_blocks:
            px1, py1, px2, py2 = b
            center_y = (py1 + py2) // 2
            card_height = MIN_CARD_HEIGHT
            card_top = max(top_limit, center_y - int(card_height * 0.7))
            card_bottom = card_top + card_height
            if card_bottom > bottom_limit:
                card_bottom = bottom_limit
                card_top = max(top_limit, card_bottom - card_height)
            card_left = 0
            card_right = w
            cards.append([card_left, card_top, card_right, card_bottom])
        cards.sort(key=lambda b: b[1])
        return cards

    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 限制检测范围
    top_limit = int(h * top_ratio)
    bottom_limit = int(h * (1 - bottom_ratio))
    logger.info(f"CV检测卡片: h={h}, w={w}, top_limit={top_limit}, bottom_limit={bottom_limit}")
    
    # 使用自适应阈值
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
    
    # 闭运算：连接断开的边缘
    # 减小水平核宽度，增加一点垂直连接，以适应卡片边缘不连续的情况
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (w // 4, 3))
    closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
    
    # 寻找轮廓
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    cards = []
    # 更严格的最小宽度要求：列表卡片通常接近整宽
    min_card_width = int(w * 0.8)
    
    temp_boxes = []
    for cnt in contours:
        x, y, cw, ch = cv2.boundingRect(cnt)
        center_y = y + ch // 2
        
        # 绿色营销券区域排除
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        roi = hsv[max(0,y):min(h,y+ch), max(0,x):min(w,x+cw)]
        green_mask = cv2.inRange(roi, np.array([35, 80, 80]), np.array([85, 255, 255]))
        green_ratio = float(cv2.countNonZero(green_mask)) / (roi.shape[0]*roi.shape[1]) if roi.size > 0 else 0.0
        
        # 圆角矩形判断：面积占比与顶点数量
        rect_area = cw * ch
        cnt_area = cv2.contourArea(cnt)
        extent = cnt_area / rect_area if rect_area > 0 else 0
        approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
        
        ok_width = cw >= min_card_width
        ok_height = ch > MIN_CARD_HEIGHT * 0.8
        ok_vertical = center_y >= int(h * 0.58) and y > top_limit and center_y < bottom_limit
        ok_shape = (len(approx) >= 6 or extent > 0.85)
        ok_color = green_ratio < 0.25
        
        if ok_width and ok_height and ok_vertical and ok_shape and ok_color:
            logger.info(f"✅ 找到候选卡片: y={y}, h={ch}, w={cw}, extent={extent:.2f}, verts={len(approx)}, green={green_ratio:.2f}")
            temp_boxes.append((x, y, cw, ch))
        else:
            if ok_width:
                logger.debug(f"❌ 排除候选: y={y}, h={ch}, w={cw}, extent={extent:.2f}, verts={len(approx)}, green={green_ratio:.2f} (limit: {top_limit}-{bottom_limit})")
    
    if not temp_boxes:
        logger.warning(f"⚠️ CV 未能在指定范围 ({top_limit}-{bottom_limit}) 内检测到任何场站卡片")
            
    # 按 Y 轴排序
    temp_boxes.sort(key=lambda b: b[1])
    
    # 再次过滤和去重，并根据高度拆分为多张卡片
    for i, box in enumerate(temp_boxes):
        x, y, cw, ch = box
        
        # 检查是否与已有的框重叠（按顶部 y 去重）
        is_duplicate = False
        for v in cards:
            if abs(y - v[1]) < 100:
                is_duplicate = True
                break
        if is_duplicate:
            continue

        padding = 2
        x1 = max(0, x + padding)
        x2 = min(w, x + cw - padding)
        y1 = max(top_limit, y + padding)
        y2 = min(bottom_limit, y + ch - padding)

        if y2 <= y1:
            continue

        effective_h = y2 - y1
        estimated_count = max(1, int(round(effective_h / float(MIN_CARD_HEIGHT))))

        if estimated_count <= 1:
            cards.append([x1, y1, x2, y2])
        else:
            seg_h = effective_h / float(estimated_count)
            for k in range(estimated_count):
                sy1 = int(round(y1 + k * seg_h))
                sy2 = int(round(y1 + (k + 1) * seg_h))
                if sy2 <= sy1:
                    continue
                if sy2 - sy1 < MIN_CARD_HEIGHT * 0.6:
                    continue
                cards.append([x1, sy1, x2, sy2])
            
    return cards

def detect_wide_rounded_card_cv(image_path, min_width_ratio=0.8, min_y_ratio=0.5):
    img = read_image(image_path)
    if img is None:
        return []
    h, w = img.shape[:2]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (7, 7), 0)
    edges = cv2.Canny(blur, 60, 180)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    results = []
    for cnt in contours:
        x, y, cw, ch = cv2.boundingRect(cnt)
        center_y = y + ch // 2
        if cw < int(w * min_width_ratio):
            continue
        if center_y < int(h * min_y_ratio):
            continue
        rect_area = cw * ch
        cnt_area = cv2.contourArea(cnt)
        if rect_area <= 0:
            continue
        extent = cnt_area / rect_area
        approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
        if extent < 0.7:
            continue
        if len(approx) >= 6 or (len(approx) == 4 and extent > 0.85):
            results.append([x, y, x + cw, y + ch])
    results.sort(key=lambda b: b[1])
    return results

def draw_rectangles(image_path, bboxes=None, click_points=None):
    """
    使用 OpenCV 在图片上绘制矩形框和点击点，生成 _vl.jpg 和 _flag.jpg
    - _vl.jpg: 仅包含矩形框，供视觉模型参考
    - _flag.jpg: 包含矩形框和点击点，供人工调试
    """
    try:
        DEBUG_BOX_COLOR = (0, 255, 0)      # 绿色矩形
        DEBUG_POINT_COLOR = (0, 0, 255)    # 红色点击点
        DEBUG_BOX_THICKNESS = 3
        
        img = read_image(image_path)
        if img is None:
            return image_path
            
        # 1. 先绘制矩形框 (VL 版)
        vl_img = img.copy()
        if bboxes:
            for box in bboxes:
                if len(box) == 4:
                    cv2.rectangle(vl_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), DEBUG_BOX_COLOR, DEBUG_BOX_THICKNESS)
        
        vl_path = image_path.replace(".jpg", "_vl.jpg")
        save_image(vl_path, vl_img)
        
        # 2. 在 VL 版基础上绘制点击点 (Flag 版)
        flag_img = vl_img.copy()
        if click_points:
            for p in click_points:
                if len(p) == 2:
                    center = (int(p[0]), int(p[1]))
                    # 绘制一个实心圆和中心点
                    cv2.circle(flag_img, center, 12, DEBUG_POINT_COLOR, -1)
                    cv2.circle(flag_img, center, 2, (255, 255, 255), -1)
                    
        flag_path = image_path.replace(".jpg", "_flag.jpg")
        save_image(flag_path, flag_img)
        
        return vl_path, flag_path
    except Exception as e:
        logger.error(f"绘制诊断图片失败: {e}")
        return image_path, 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}")

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. 获取供应商代号 (如 TelaiDian)
    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 (如 TelaiDian.T4_TelaiDian)
    if name == supplier_code:
        return parent_logger
    return logging.getLogger(f"{supplier_code}.{name}")