import logging import os import cv2 import numpy as np import time import json import hashlib from Apps.AiTeJiYiChong.Config.Setting import BOTTOM_SAFE_EXCLUDE_RATIO from Config.Config import TEMP_IMAGE_DIR logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) def clean_station_name(name): """ 清理场站名称,去掉末尾的省略号和空格 """ if not name: return "" # 去掉末尾的 ... 或 …… name = name.strip() while name.endswith(".") or name.endswith("。"): name = name[:-1] return name.strip() 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 # 使用 cv2.imencode 转为 jpg 字节流,避免原始 numpy 数组可能的细微差异(虽然通常没问题) 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: # 使用 np.fromfile 解决中文路径问题 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 take_screenshot(d, image_uuid, save_dir=TEMP_IMAGE_DIR): path = os.path.join(save_dir, f"{image_uuid}.jpg") os.makedirs(save_dir, exist_ok=True) d.screenshot(path) return path def clear_temp_dir(save_dir=TEMP_IMAGE_DIR): """清空临时目录中的所有文件""" if not os.path.exists(save_dir): return logger.info(f"正在清空临时目录: {save_dir}") for file in os.listdir(save_dir): file_path = os.path.join(save_dir, file) try: if os.path.isfile(file_path): os.remove(file_path) elif os.path.isdir(file_path): import shutil shutil.rmtree(file_path) except Exception as e: logger.error(f"无法删除文件 {file_path}: {e}") def click_image_template(d, template_path, timeout=5.0, threshold=0.8): """ 使用 OpenCV 模板匹配查找并点击图片 """ if not os.path.exists(template_path): logger.info(f"Template file not found: {template_path}") return False template = read_image(template_path) if template is None: logger.info(f"Failed to load template: {template_path}") return False t_h, t_w = template.shape[:2] start_time = time.time() while time.time() - start_time < timeout: temp_uuid = "temp_click_check" screenshot_path = take_screenshot(d, temp_uuid, save_dir=TEMP_IMAGE_DIR) target = read_image(screenshot_path) if target is None: time.sleep(0.5) continue # 多尺度匹配 best_match = None for scale in np.linspace(0.8, 1.2, 5): resized = cv2.resize(template, (int(t_w * scale), int(t_h * scale))) res = cv2.matchTemplate(target, resized, cv2.TM_CCOEFF_NORMED) _, max_val, _, max_loc = cv2.minMaxLoc(res) if best_match is None or max_val > best_match[0]: best_match = (max_val, max_loc, resized.shape[1], resized.shape[0]) if best_match and best_match[0] >= threshold: max_val, max_loc, r_w, r_h = best_match center_x = max_loc[0] + r_w // 2 center_y = max_loc[1] + r_h // 2 logger.info(f"成功点击图片模板: {template_path}, 匹配度: {max_val:.2f}") d.click(center_x, center_y) return True time.sleep(0.5) return False 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. 获取供应商代号 (如 AiTeJiYiChong) 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}") def find_template_coords(img_path, template_path, threshold=0.8): """ 在图片中查找模板并返回中心坐标 """ if not os.path.exists(img_path) or not os.path.exists(template_path): return None img = read_image(img_path) template = read_image(template_path) if img is None or template is None: return None t_h, t_w = template.shape[:2] # 多尺度匹配 best_match = None for scale in np.linspace(0.8, 1.2, 5): resized = cv2.resize(template, (int(t_w * scale), int(t_h * scale))) if resized.shape[0] > img.shape[0] or resized.shape[1] > img.shape[1]: continue res = cv2.matchTemplate(img, resized, cv2.TM_CCOEFF_NORMED) _, max_val, _, max_loc = cv2.minMaxLoc(res) if best_match is None or max_val > best_match[0]: best_match = (max_val, max_loc, resized.shape[1], resized.shape[0]) if best_match and best_match[0] >= threshold: max_val, max_loc, r_w, r_h = best_match center_x = max_loc[0] + r_w // 2 center_y = max_loc[1] + r_h // 2 return (center_x, center_y, max_val) return None def crop_cards_from_image(img_path, output_dir=None, save_debug=True): """ 从图片中裁剪场站卡片并生成 _flag.jpg 和 _vl.jpg 算法:以导航图标 (arrow.jpg) 为主要锚点,辅以红色价格区域,向上/下/左/右探测背景边界 """ logger.info(f"Processing: {img_path}") if not os.path.exists(img_path): return [] img = read_image(img_path) if img is None: return [] h, w = img.shape[:2] hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) s_channel = hsv[:, :, 1] anchors = [] # 1. 使用导航图标 (arrow.jpg) 进行模板匹配 template_path = os.path.join(os.path.dirname(__file__), "BiaoShi", "arrow.jpg") if os.path.exists(template_path): template = read_image(template_path) if template is not None: t_h, t_w = template.shape[:2] # 模板匹配 res = cv2.matchTemplate(img, template, cv2.TM_CCOEFF_NORMED) threshold = 0.7 loc = np.where(res >= threshold) # 去重并记录锚点 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"Found {len(matched_points)} anchors via arrow template matching.") # 2. 如果模板匹配找得不够,或者作为补充,使用红色价格区域 # 进一步放宽红色范围 lower_red1 = np.array([0, 30, 30]) upper_red1 = np.array([15, 255, 255]) lower_red2 = np.array([150, 30, 30]) upper_red2 = np.array([180, 255, 255]) mask1 = cv2.inRange(hsv, lower_red1, upper_red1) mask2 = cv2.inRange(hsv, lower_red2, upper_red2) red_mask = cv2.bitwise_or(mask1, mask2) kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 20)) morphed = cv2.dilate(red_mask, kernel, iterations=1) contours, _ = cv2.findContours(morphed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) red_anchors_count = 0 red_rects = [] # 记录所有红色的矩形区域 for cnt in contours: x, y, cw, ch = cv2.boundingRect(cnt) if 20 < cw < 600 and 10 < ch < 150 and y > h * 0.2: # 在原始 red_mask 中精确查找该区域的边界,避免膨胀带来的误差 roi_mask = red_mask[y:y+ch, x:x+cw] points = cv2.findNonZero(roi_mask) if points is not None: rx, ry, rcw, rch = cv2.boundingRect(points) # 转换回原图坐标 exact_rect = (x + rx, y + ry, rcw, rch) red_rects.append(exact_rect) ax, ay = exact_rect[0] + exact_rect[2] // 2, exact_rect[1] + exact_rect[3] // 2 else: red_rects.append((x, y, cw, ch)) ax, ay = x + cw // 2, y + ch // 2 # 检查是否与已有锚点重合 is_duplicate = False for ex, ey in anchors: if abs(ey - ay) < 100: is_duplicate = True break if not is_duplicate: anchors.append((ax, ay)) red_anchors_count += 1 logger.info(f"Found {red_anchors_count} additional anchors via red price detection.") # 3. 定位背景分隔行 (Blue Separator Rows) # 计算每一行的平均饱和度和色调 h_channel = hsv[:, :, 0] row_s_means = np.mean(s_channel[:, int(w*0.2):int(w*0.8)], axis=1) row_h_means = np.mean(h_channel[:, int(w*0.2):int(w*0.8)], axis=1) row_s_stds = np.std(s_channel[:, int(w*0.2):int(w*0.8)], axis=1) # 分隔行的特征:浅蓝色 (H~105, S>15) 且整行均匀 (std小) is_separator = (row_s_means > 15) & (row_h_means > 90) & (row_h_means < 120) & (row_s_stds < 10) # 找出所有“非分隔”区域 segments = [] start_y = None for y in range(h): if not is_separator[y]: if start_y is None: start_y = y else: if start_y is not None: segments.append((start_y, y)) start_y = None if start_y is not None: segments.append((start_y, h)) # 合并非常接近的区域 (可能被误认为分隔线的行) merged_segments = [] if segments: curr_start, curr_end = segments[0] for i in range(1, len(segments)): next_start, next_end = segments[i] if next_start - curr_end < 20: # 间隙小于 20 像素则合并 curr_end = next_end else: merged_segments.append((curr_start, curr_end)) curr_start, curr_end = next_start, next_end merged_segments.append((curr_start, curr_end)) logger.info(f"Found {len(merged_segments)} merged segments: {merged_segments}") final_cards = [] for b_start, b_end in merged_segments: card_h = b_end - b_start # 场站卡片高度通常在 250 到 500 之间 if card_h > 150: # 检查这个段落里是否有锚点 segment_anchors = [ (ax, ay) for ax, ay in anchors if b_start - 20 < ay < b_end + 20 ] if segment_anchors: # 优化下边界:如果存在红色价格,以下边界为准 (用户建议) segment_red_rects = [ (rx, ry, rcw, rch) for rx, ry, rcw, rch in red_rects if b_start < ry < b_end ] y2_refined = b_end if segment_red_rects: # 找到最下方的红色区域 max_red_bottom = max([ry + rch for rx, ry, rcw, rch in segment_red_rects]) # 用户建议:发现红色字结束就停止计算下边界。给予微小缓冲空间 (5px) y2_refined = min(b_end, max_red_bottom + 5) logger.info(f"Refined y2 from {b_end} to {y2_refined} based on red text.") # 如果段落太大,可能包含了多个卡片(分隔线没断开) if card_h > 600: logger.info(f"Segment at y=[{b_start}, {b_end}] is too large ({card_h}), attempting split by anchors...") segment_anchors.sort(key=lambda a: a[1]) # 只有当锚点之间距离足够大时才拆分 splits = [b_start] for i in range(len(segment_anchors) - 1): ay1 = segment_anchors[i][1] ay2 = segment_anchors[i+1][1] if ay2 - ay1 > 200: # 锚点间距大于 200 才考虑拆分 # 在两个锚点之间找最像分隔线的行 (饱和度最高) split_y = ay1 + np.argmax(row_s_means[ay1:ay2]) splits.append(split_y) splits.append(b_end) for i in range(len(splits) - 1): s1, s2 = splits[i], splits[i+1] if s2 - s1 > 150: # 对拆分后的每个部分也尝试优化下边界 part_red_rects = [ (rx, ry, rcw, rch) for rx, ry, rcw, rch in red_rects if s1 < ry < s2 ] s2_refined = s2 if part_red_rects: max_red_bottom = max([ry + rch for rx, ry, rcw, rch in part_red_rects]) s2_refined = min(s2, max_red_bottom + 5) final_cards.append((s1, s2_refined, int(w*0.02), int(w*0.98))) logger.info(f"Added split card: y=[{s1}, {s2_refined}]") else: final_cards.append((b_start, y2_refined, int(w*0.02), int(w*0.98))) logger.info(f"Added card: y=[{b_start}, {y2_refined}], original_h={card_h}") # 4. 排序 (按 y1 从上到下) final_cards.sort(key=lambda c: c[0]) # 保存结果图 if output_dir is None: output_dir = os.path.dirname(img_path) base_name = os.path.basename(img_path) stem, ext = os.path.splitext(base_name) debug_img = img.copy() # _flag.jpg vl_img = img.copy() # _vl.jpg json_data = {"image": base_name, "width": w, "height": h, "cards": []} for idx, (y1, y2, x1, x2) in enumerate(final_cards): # 计算点击点 (卡片上方区域) click_x = int(x1 + (x2 - x1) * 0.2) click_y = int(y1 + (y2 - y1) * 0.2) # 在 flag 图上画绿框和红点 cv2.rectangle(debug_img, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.circle(debug_img, (click_x, click_y), 10, (0, 0, 255), -1) # 在 vl 图上只画绿框 cv2.rectangle(vl_img, (x1, y1), (x2, y2), (0, 255, 0), 2) json_data["cards"].append({ "id": idx + 1, "rect": [int(x1), int(y1), int(x2), int(y2)], "bounds_norm": { "left": round(float(x1) / w, 4), "top": round(float(y1) / h, 4), "right": round(float(x2) / w, 4), "bottom": round(float(y2) / h, 4) }, "click_point": [int(click_x), int(click_y)] }) # 保存文件 if save_debug: save_image(os.path.join(output_dir, f"{stem}_flag{ext}"), debug_img) save_image(os.path.join(output_dir, f"{stem}_vl{ext}"), vl_img) with open(os.path.join(output_dir, f"{stem}.json"), 'w', encoding='utf-8') as f: json.dump(json_data, f, ensure_ascii=False, indent=2) logger.info(f"Generated _flag and _vl images for {len(final_cards)} cards.") return json_data