740 lines
25 KiB
Python
740 lines
25 KiB
Python
import logging
|
||
import os
|
||
import cv2
|
||
import numpy as np
|
||
import time
|
||
import hashlib
|
||
from Apps.XinDianTu.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 read_image(path):
|
||
"""读取图片，支持中文路径及鲁棒性检查"""
|
||
try:
|
||
if not path or not os.path.exists(path):
|
||
return None
|
||
if os.path.getsize(path) == 0:
|
||
return None
|
||
return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
|
||
except Exception as e:
|
||
logger.error(f"读取图片失败 {path}: {e}")
|
||
return None
|
||
|
||
|
||
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 save_image(path, img):
|
||
"""保存图片，支持中文路径及鲁棒性检查"""
|
||
try:
|
||
if img is None:
|
||
return False
|
||
|
||
# 确保目录存在
|
||
dir_name = os.path.dirname(path)
|
||
if dir_name and not os.path.exists(dir_name):
|
||
os.makedirs(dir_name, exist_ok=True)
|
||
|
||
ext = os.path.splitext(path)[1]
|
||
if not ext:
|
||
ext = ".jpg"
|
||
|
||
success, encoded_img = cv2.imencode(ext, img)
|
||
if success:
|
||
encoded_img.tofile(path)
|
||
return True
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"保存图片失败 {path}: {e}")
|
||
return False
|
||
|
||
|
||
# 截图
|
||
def take_screenshot(d, image_uuid, save_dir=TEMP_IMAGE_DIR):
|
||
path = f"{save_dir}/{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 is_background_dimmed(image_path, threshold=80):
|
||
"""
|
||
检测背景是否被暗色蒙板覆盖 (Image Mask Detection)
|
||
原理：计算屏幕四周边缘区域的平均亮度。如果边缘区域普遍较暗（蒙板效果），则返回 True。
|
||
:param image_path: 截图路径
|
||
:param threshold: 亮度阈值，低于此值认为是蒙板 (0-255)
|
||
:return: bool
|
||
"""
|
||
if not os.path.exists(image_path):
|
||
return False
|
||
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return False
|
||
|
||
h, w = img.shape[:2]
|
||
# 转换为灰度图计算亮度
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 定义边缘区域 (四周各取 5% 的宽度/高度)
|
||
edge_h = int(h * 0.05)
|
||
edge_w = int(w * 0.05)
|
||
|
||
# 提取四个角落/边缘块
|
||
top_edge = gray[0:edge_h, :]
|
||
bottom_edge = gray[h-edge_h:h, :]
|
||
left_edge = gray[:, 0:edge_w]
|
||
right_edge = gray[:, w-edge_w:w]
|
||
|
||
# 计算平均亮度
|
||
avg_brightness = (np.mean(top_edge) + np.mean(bottom_edge) + np.mean(left_edge) + np.mean(right_edge)) / 4
|
||
|
||
logger.info(f"Background Dimmed Check: Avg Brightness = {avg_brightness:.2f} (Threshold: {threshold})")
|
||
|
||
# 如果平均亮度低于阈值，说明背景被压暗了
|
||
return avg_brightness < threshold
|
||
|
||
|
||
def detect_black_agree_button(image_path, debug_dir=None):
|
||
"""
|
||
通过计算机图形学检测黑色的"同意"按钮 (Image 1 场景)
|
||
特征：黑色圆角矩形，位于屏幕中下部，面积适中
|
||
:param image_path: 截图路径
|
||
:return: (x, y) 坐标中心点，如果未找到返回 None
|
||
"""
|
||
if not os.path.exists(image_path):
|
||
return None
|
||
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return None
|
||
|
||
h, w = img.shape[:2]
|
||
|
||
# 转换为HSV颜色空间，因为黑色更容易过滤
|
||
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
||
|
||
# 定义黑色的范围
|
||
# 黑色：V (Brightness) 很低
|
||
lower_black = np.array([0, 0, 0])
|
||
upper_black = np.array([180, 255, 40]) # V < 40 认为是黑色
|
||
|
||
mask = cv2.inRange(hsv, lower_black, upper_black)
|
||
|
||
# 限制搜索区域：通常在屏幕下半部分
|
||
roi_top = int(h * 0.4)
|
||
roi_bottom = int(h * 0.8)
|
||
roi_mask = np.zeros_like(mask)
|
||
roi_mask[roi_top:roi_bottom, :] = mask[roi_top:roi_bottom, :]
|
||
|
||
# 形态学操作：去除噪点，连接断开的区域
|
||
kernel = np.ones((5, 5), np.uint8)
|
||
roi_mask = cv2.morphologyEx(roi_mask, cv2.MORPH_CLOSE, kernel)
|
||
roi_mask = cv2.morphologyEx(roi_mask, cv2.MORPH_OPEN, kernel)
|
||
|
||
# 查找轮廓
|
||
contours, _ = cv2.findContours(roi_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
best_cnt = None
|
||
max_area = 0
|
||
|
||
for cnt in contours:
|
||
area = cv2.contourArea(cnt)
|
||
x, y, cw, ch = cv2.boundingRect(cnt)
|
||
aspect_ratio = float(cw) / ch
|
||
|
||
# 过滤条件
|
||
# 1. 面积要够大 (例如 > 屏幕面积的 1%)
|
||
if area < (w * h * 0.01):
|
||
continue
|
||
|
||
# 2. 宽高比：通常按钮是扁长条，例如 > 2.0
|
||
if aspect_ratio < 2.0 or aspect_ratio > 10.0:
|
||
continue
|
||
|
||
# 3. 宽度：通常占据屏幕宽度的 50% 以上
|
||
if cw < (w * 0.5):
|
||
continue
|
||
|
||
if area > max_area:
|
||
max_area = area
|
||
best_cnt = cnt
|
||
|
||
if best_cnt is not None:
|
||
x, y, cw, ch = cv2.boundingRect(best_cnt)
|
||
center_x = x + cw // 2
|
||
center_y = y + ch // 2
|
||
logger.info(f"Found Black Agree Button at ({center_x}, {center_y}), Size: {cw}x{ch}")
|
||
|
||
if debug_dir:
|
||
os.makedirs(debug_dir, exist_ok=True)
|
||
debug_img = img.copy()
|
||
cv2.rectangle(debug_img, (x, y), (x + cw, y + ch), (0, 0, 255), 2)
|
||
cv2.circle(debug_img, (center_x, center_y), 5, (0, 255, 0), -1)
|
||
save_image(os.path.join(debug_dir, "debug_agree_btn.jpg"), debug_img)
|
||
|
||
return (center_x, center_y)
|
||
|
||
return None
|
||
|
||
|
||
|
||
|
||
|
||
|
||
def find_expand_button_position(image_path, debug_dir=None, debug_filename_prefix=None):
|
||
"""
|
||
通过几何特征识别"全部时段"按钮的位置
|
||
特征：该行左侧(30%)和右侧(30%)基本为空白，中间有内容
|
||
:param image_path: 截图路径
|
||
:param debug_dir: 调试图片保存目录，如果为None则不保存
|
||
:param debug_filename_prefix: 调试图片文件名前缀
|
||
:return: (x, y) 坐标中心点，如果未找到返回 None
|
||
"""
|
||
if not os.path.exists(image_path):
|
||
return None
|
||
|
||
img = read_image(image_path)
|
||
if img is None:
|
||
return None
|
||
|
||
h, w = img.shape[:2]
|
||
|
||
# 转灰度
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 边缘检测
|
||
edges = cv2.Canny(gray, 50, 150)
|
||
|
||
# 定义区域
|
||
left_w = int(w * 0.35)
|
||
right_w = int(w * 0.65)
|
||
|
||
# 提取各部分边缘
|
||
left_part = edges[:, :left_w]
|
||
right_part = edges[:, right_w:]
|
||
center_part = edges[:, left_w:right_w]
|
||
|
||
# 计算每一行的边缘点数量
|
||
row_sum_left = np.sum(left_part, axis=1) / 255.0
|
||
row_sum_right = np.sum(right_part, axis=1) / 255.0
|
||
row_sum_center = np.sum(center_part, axis=1) / 255.0
|
||
|
||
# 筛选条件：左右边缘点很少，中间边缘点较多
|
||
# 阈值可以根据实际情况调整
|
||
# 允许少量噪点，所以不是严格的0，而是小于某个较小值（例如宽度的1%）
|
||
noise_threshold = 2 # 允许2个像素的噪点
|
||
content_threshold = 5 # 中间至少有5个像素的边缘
|
||
|
||
candidates = (row_sum_left <= noise_threshold) & \
|
||
(row_sum_right <= noise_threshold) & \
|
||
(row_sum_center >= content_threshold)
|
||
|
||
# 找到连续的候选行
|
||
y_indices = np.where(candidates)[0]
|
||
|
||
if len(y_indices) == 0:
|
||
return None
|
||
|
||
# 将连续行分组
|
||
segments = []
|
||
if len(y_indices) > 0:
|
||
start = y_indices[0]
|
||
prev = y_indices[0]
|
||
for y in y_indices[1:]:
|
||
if y > prev + 5: # 允许5像素断裂
|
||
segments.append((start, prev))
|
||
start = y
|
||
prev = y
|
||
segments.append((start, prev))
|
||
|
||
# 筛选最合适的段
|
||
best_segment = None
|
||
|
||
# 我们期望按钮在屏幕中下部，且高度适中（例如 20-100px）
|
||
# 且通常是在价格表下方。假设价格表占据了屏幕上部。
|
||
# 我们可以简单地取符合条件的段中，Y值最大的那个（最靠下的），或者最符合"中间有字"特征的。
|
||
# 考虑到页面底部可能有其他干扰，取"中下部"的一个。
|
||
|
||
valid_segments = []
|
||
for start, end in segments:
|
||
height = end - start
|
||
mid_y = (start + end) // 2
|
||
|
||
# 过滤掉太高或太矮的区域
|
||
if height < 20 or height > 150:
|
||
continue
|
||
|
||
# 过滤掉顶部的区域（可能是标题栏误判）
|
||
if mid_y < h * 0.3:
|
||
continue
|
||
|
||
# 过滤掉底部的区域（可能是底部按钮）
|
||
if mid_y > h * 0.9:
|
||
continue
|
||
|
||
valid_segments.append((start, end))
|
||
|
||
if not valid_segments:
|
||
return None
|
||
|
||
# 如果有多个，通常"全部时段"是在价格表之后，紧接着的一个
|
||
# 这里我们取第一个（最靠上的）符合条件的段，因为它紧跟在价格表下方
|
||
# 或者取所有段中，中间内容最"紧凑"的？
|
||
# 让我们简单点，取第一个符合条件的段。
|
||
best_segment = valid_segments[0]
|
||
|
||
start, end = best_segment
|
||
center_y = (start + end) // 2
|
||
center_x = w // 2
|
||
|
||
if debug_dir:
|
||
os.makedirs(debug_dir, exist_ok=True)
|
||
debug_img = img.copy()
|
||
# 画出识别区域
|
||
cv2.rectangle(debug_img, (0, start), (w, end), (0, 255, 0), 2)
|
||
# 画出红点
|
||
cv2.circle(debug_img, (center_x, center_y), 10, (0, 0, 255), -1)
|
||
|
||
# 保存 flag 图片
|
||
if debug_filename_prefix:
|
||
# Sanitize filename: remove invalid chars
|
||
import re
|
||
safe_prefix = re.sub(r'[\\/*?:"<>|]', '_', str(debug_filename_prefix))
|
||
debug_name = f"{safe_prefix}_flag_expand.jpg"
|
||
else:
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
debug_name = f"{timestamp}_flag_expand.jpg"
|
||
|
||
debug_path = os.path.join(debug_dir, debug_name)
|
||
save_image(debug_path, debug_img)
|
||
logger.info(f"Saved debug image to {debug_path}")
|
||
|
||
return (int(center_x), int(center_y))
|
||
|
||
|
||
def get_row_stats(gray):
|
||
"""
|
||
计算每一行的统计特征
|
||
"""
|
||
h, w = gray.shape
|
||
|
||
# 中央区域 (用于检测内容)
|
||
center_x = w // 2
|
||
strip_w = 100
|
||
center_strip = gray[:, center_x - 50: center_x + 50]
|
||
|
||
# 边缘区域 (用于检测背景/边距)
|
||
# 假设边距至少有 10px
|
||
edge_strip = gray[:, 0:20]
|
||
|
||
row_means = np.mean(center_strip, axis=1)
|
||
row_stds = np.std(center_strip, axis=1)
|
||
edge_means = np.mean(edge_strip, axis=1)
|
||
|
||
return row_means, row_stds, edge_means
|
||
|
||
|
||
def clean_station_name(name):
|
||
"""
|
||
清理场站名称，移除结尾的省略号及多余空格，方便比对
|
||
"""
|
||
if not name: return ""
|
||
name = name.strip()
|
||
# 移除结尾的 . (通常是截断标识)
|
||
while name.endswith("."):
|
||
name = name[:-1]
|
||
return name.strip()
|
||
|
||
|
||
def get_file_md5(path):
|
||
"""计算文件的 MD5 值"""
|
||
if not os.path.exists(path):
|
||
return ""
|
||
hash_md5 = hashlib.md5()
|
||
with open(path, "rb") as f:
|
||
for chunk in iter(lambda: f.read(4096), b""):
|
||
hash_md5.update(chunk)
|
||
return hash_md5.hexdigest()
|
||
|
||
|
||
def crop_cards_from_image(img_path, output_dir=None, save_debug=True):
|
||
"""
|
||
从图片中裁剪场站卡片
|
||
:param img_path: 图片路径
|
||
:param output_dir: 输出目录，默认与 img_path 相同
|
||
:param save_debug: 是否保存调试图 (_flag.jpg)
|
||
:return: 裁剪出的卡片列表，每项包含 (out_path, (click_x, click_y))
|
||
"""
|
||
logger.info(f"Processing: {img_path}")
|
||
if not os.path.exists(img_path):
|
||
logger.info(f"Error: File not found {img_path}")
|
||
return []
|
||
|
||
img = read_image(img_path)
|
||
if img is None:
|
||
logger.info(f"Error: Failed to load image {img_path}")
|
||
return []
|
||
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
h, w = gray.shape
|
||
|
||
row_means, row_stds, edge_means = get_row_stats(gray)
|
||
|
||
# 参数定义
|
||
# 背景灰度值范围 (根据 Analyze2.py 的输出，背景约 242)
|
||
BG_GRAY_MIN = 230
|
||
BG_GRAY_MAX = 250
|
||
|
||
# 间隙判定：中央区域也是均匀的灰度
|
||
# 真正的背景灰度约 242。卡片内的伪背景(不够白)约 246-248。
|
||
# 所以降低 GAP_MEAN_MAX 以区分 Gap 和 Dirty White。
|
||
GAP_MEAN_MIN = 235
|
||
GAP_MEAN_MAX = 244
|
||
GAP_STD_MAX = 10.0
|
||
|
||
segments = []
|
||
is_in_card = False
|
||
start_y = 0
|
||
|
||
# 记录原始的行状态，用于后续可能的合并优化
|
||
row_is_card = []
|
||
|
||
for y in range(h):
|
||
rm = row_means[y]
|
||
rs = row_stds[y]
|
||
em = edge_means[y]
|
||
|
||
# 1. 判定当前行是否可能是卡片的一部分
|
||
# 条件A: 边缘是灰色的 (说明有边距，排除了全宽的Header/Footer)
|
||
has_gray_margin = (BG_GRAY_MIN <= em <= BG_GRAY_MAX)
|
||
|
||
# 条件B: 中央不是均匀的背景灰 (说明有内容，或者是白色底)
|
||
# 如果是 Gap，则 Mean 在 Gap范围 且 Std 很小
|
||
is_gap = (GAP_MEAN_MIN <= rm <= GAP_MEAN_MAX) and (rs < GAP_STD_MAX)
|
||
|
||
is_card_row = has_gray_margin and not is_gap
|
||
row_is_card.append(is_card_row)
|
||
|
||
# 简单的形态学闭运算：填补卡片内部的小裂缝
|
||
# 如果一个 False (Gap) 的上下都是 True (Card)，且 Gap 长度很短，则将其视为 Card
|
||
# 允许的裂缝最大长度
|
||
# [优化] 减小填补阈值，避免把两个卡片中间的真 Gap 填补了导致粘连
|
||
# 之前是 10，现在改为 3。
|
||
MAX_GAP_FILL = 3
|
||
|
||
cleaned_row_is_card = row_is_card[:]
|
||
|
||
# 查找连续的 False 区域
|
||
i = 0
|
||
while i < h:
|
||
if not cleaned_row_is_card[i]:
|
||
# Found a gap start
|
||
gap_start = i
|
||
while i < h and not cleaned_row_is_card[i]:
|
||
i += 1
|
||
gap_end = i
|
||
gap_len = gap_end - gap_start
|
||
|
||
# Check context
|
||
prev_is_card = (gap_start > 0) and cleaned_row_is_card[gap_start - 1]
|
||
next_is_card = (gap_end < h) and cleaned_row_is_card[gap_end]
|
||
|
||
if prev_is_card and next_is_card and gap_len <= MAX_GAP_FILL:
|
||
# Fill the gap
|
||
for k in range(gap_start, gap_end):
|
||
cleaned_row_is_card[k] = True
|
||
else:
|
||
i += 1
|
||
|
||
# 根据 cleaned_row_is_card 生成 segments
|
||
is_in_card = False
|
||
start_y = 0
|
||
for y in range(h):
|
||
if cleaned_row_is_card[y]:
|
||
if not is_in_card:
|
||
is_in_card = True
|
||
start_y = y
|
||
else:
|
||
if is_in_card:
|
||
is_in_card = False
|
||
end_y = y
|
||
height = end_y - start_y
|
||
|
||
# 过滤太矮的区域
|
||
if height > 100: # 稍微放宽一点，150 -> 100
|
||
segments.append((start_y, end_y))
|
||
|
||
# 处理最后一个 segment
|
||
if is_in_card:
|
||
end_y = h
|
||
height = end_y - start_y
|
||
if height > 100:
|
||
segments.append((start_y, end_y))
|
||
|
||
# [新增] 触底过滤：过滤掉延伸到图片底部的 Segment
|
||
# 底部通常是导航栏或被截断的卡片
|
||
BOTTOM_MARGIN = 50 # 距离底部多少像素内算触底
|
||
valid_segments_scan = []
|
||
for y1, y2 in segments:
|
||
if y2 >= h - BOTTOM_MARGIN:
|
||
logger.info(f" Dropping segment Y={y1}-{y2} because it touches the bottom (H={h}).")
|
||
continue
|
||
valid_segments_scan.append((y1, y2))
|
||
segments = valid_segments_scan
|
||
|
||
logger.info(f" Found {len(segments)} potential segments based on vertical scan.")
|
||
|
||
# 1.5 Group Segments by Large Gaps (to separate Header / List / Footer)
|
||
# The filter bar usually creates a large gap (e.g. > 60px).
|
||
SECTION_GAP_MIN = 60
|
||
groups = []
|
||
if segments:
|
||
current_group = [segments[0]]
|
||
for i in range(1, len(segments)):
|
||
prev_end = segments[i - 1][1]
|
||
curr_start = segments[i][0]
|
||
gap = curr_start - prev_end
|
||
|
||
if gap > SECTION_GAP_MIN:
|
||
groups.append(current_group)
|
||
current_group = []
|
||
|
||
current_group.append(segments[i])
|
||
groups.append(current_group)
|
||
|
||
logger.info(f" Found {len(groups)} segment groups.")
|
||
|
||
# Select the Target Group
|
||
# Strategy: Select the LAST group that has at least one 'substantial' segment (H > 150).
|
||
# This assumes the Station List is the main content and usually at the bottom (scrolling area).
|
||
target_group = []
|
||
for g in reversed(groups):
|
||
# Check if group has substantial content
|
||
has_content = False
|
||
for (y1, y2) in g:
|
||
if (y2 - y1) > 150:
|
||
has_content = True
|
||
break
|
||
|
||
if has_content:
|
||
target_group = g
|
||
break
|
||
|
||
if not target_group and segments:
|
||
# Fallback to all segments if no substantial group found
|
||
target_group = segments
|
||
|
||
logger.info(f" Selected group with {len(target_group)} segments.")
|
||
|
||
# 2. 确定每个卡片的左右边界 (Width Refinement)
|
||
# 收集所有 segment 的建议边界，然后统一
|
||
|
||
candidate_x1 = []
|
||
candidate_x2 = []
|
||
|
||
temp_valid_segments = []
|
||
|
||
for y1, y2 in target_group:
|
||
# 取中间一段行来分析
|
||
mid_y = (y1 + y2) // 2
|
||
sample_h = min(10, y2 - y1)
|
||
sample_rows = gray[mid_y - sample_h // 2: mid_y + sample_h // 2, :]
|
||
col_means = np.mean(sample_rows, axis=0)
|
||
|
||
# 从左向右找第一个非灰像素
|
||
x1 = 0
|
||
for x in range(w // 2):
|
||
if not (BG_GRAY_MIN <= col_means[x] <= BG_GRAY_MAX):
|
||
x1 = x
|
||
break
|
||
|
||
# 从右向左找第一个非灰像素
|
||
x2 = w
|
||
for x in range(w - 1, w // 2, -1):
|
||
if not (BG_GRAY_MIN <= col_means[x] <= BG_GRAY_MAX):
|
||
x2 = x
|
||
break
|
||
|
||
if x1 >= x2 or (x2 - x1) < w * 0.5:
|
||
# 如果这个 segment 无法确定宽度，可能不是有效卡片，暂不参与宽度投票
|
||
# 但为了不漏掉，暂时先记录，用默认值填充
|
||
logger.warning(f" Warning: Segment {y1}-{y2} has weird width {x2 - x1}.")
|
||
pass
|
||
else:
|
||
candidate_x1.append(x1)
|
||
candidate_x2.append(x2)
|
||
|
||
temp_valid_segments.append((y1, y2))
|
||
|
||
# 计算统一宽度
|
||
if not candidate_x1:
|
||
logger.info(" No valid width detected. Using default.")
|
||
final_x1 = 0
|
||
final_x2 = w
|
||
else:
|
||
# 使用中位数或众数来消除噪声
|
||
# 考虑到对齐，Min x1 和 Max x2 可能更合适？或者 Median。
|
||
# 通常卡片是对齐的，所以 x1 应该几乎一样。
|
||
final_x1 = int(np.median(candidate_x1))
|
||
final_x2 = int(np.median(candidate_x2))
|
||
|
||
# 稍微加点 Padding (但不要超过图片边界)
|
||
final_x1 = max(0, final_x1 - 5)
|
||
final_x2 = min(w, final_x2 + 5)
|
||
|
||
logger.info(f" Unified Width: X={final_x1}-{final_x2}, W={final_x2 - final_x1}")
|
||
|
||
# 3. 过滤高度异常的卡片 (Height Filtering)
|
||
# 计算所有潜在卡片的高度
|
||
final_cards = []
|
||
if not temp_valid_segments:
|
||
logger.info(" No segments found.")
|
||
else:
|
||
heights = [y2 - y1 for y1, y2 in temp_valid_segments]
|
||
if not heights:
|
||
logger.info(" No heights to calculate.")
|
||
else:
|
||
max_h = max(heights)
|
||
median_h = np.median(heights)
|
||
|
||
# 策略：如果高度差异较大，丢弃过小的卡片
|
||
# 阈值设定为 Median 的 75% (允许一定程度的偏差，如 220 vs 258 是允许的，但 150 vs 250 应被过滤)
|
||
# 使用 Median 比 Max 更稳健，避免被单个超高卡片(如广告)带偏
|
||
threshold_h = median_h * 0.70
|
||
|
||
for (y1, y2), h in zip(temp_valid_segments, heights):
|
||
if h < threshold_h:
|
||
logger.info(
|
||
f" Filtering out segment Y={y1}-{y2} (H={h}) because it's too short (Threshold={threshold_h:.1f}).")
|
||
else:
|
||
final_cards.append((y1, y2, final_x1, final_x2))
|
||
logger.info(f" Card: Y={y1}-{y2}, X={final_x1}-{final_x2}, H={h}")
|
||
|
||
# 4. 保存结果
|
||
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()
|
||
|
||
results = []
|
||
|
||
# 准备 JSON 数据结构
|
||
json_data = {
|
||
"image": base_name,
|
||
"width": w,
|
||
"height": h,
|
||
"cards": []
|
||
}
|
||
|
||
# 准备 _vl.jpg (只画框，不画红点)
|
||
vl_img = img.copy()
|
||
|
||
logger.info(f" Step [2.1/VL] 准备在 VL 图片上绘制 {len(final_cards)} 个场站的绿色方框...")
|
||
|
||
for idx, (y1, y2, x1, x2) in enumerate(final_cards):
|
||
# 计算点击点 (左上角，避免被底部按钮遮挡)
|
||
# 策略：X偏移 15%, Y偏移 20%
|
||
w_card = x2 - x1
|
||
h_card = y2 - y1
|
||
click_x = int(x1 + w_card * 0.15)
|
||
click_y = int(y1 + h_card * 0.20)
|
||
|
||
# [修改] 不再保存单张子图，只记录元数据
|
||
# card = img[y1:y2, x1:x2]
|
||
# 文件名添加坐标: _ClickX_ClickY
|
||
# out_name = f"{stem}_{idx + 1}_{click_x}_{click_y}{ext}"
|
||
# out_path = os.path.join(output_dir, out_name)
|
||
# cv2.imwrite(out_path, card)
|
||
# logger.info(f" Saved {out_path}")
|
||
|
||
# results.append((out_path, (click_x, click_y)))
|
||
|
||
# 在标记图上画红点 (实心圆, 半径10, 红色BGR)
|
||
cv2.circle(debug_img, (click_x, click_y), 10, (0, 0, 255), -1)
|
||
# [修改] 必须画绿框，因为后续视觉模型依赖这个框来识别范围
|
||
cv2.rectangle(debug_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
|
||
# 在 _vl 图上只画绿框
|
||
cv2.rectangle(vl_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
|
||
# 收集 JSON 数据
|
||
card_info = {
|
||
"id": idx + 1,
|
||
"rect": [x1, y1, x2, y2],
|
||
"click_point": [click_x, click_y]
|
||
}
|
||
json_data["cards"].append(card_info)
|
||
|
||
# 记录区域信息供调用者使用 (如果需要)
|
||
# 格式: (None, (click_x, click_y), (x1, y1, x2, y2))
|
||
results.append((None, (click_x, click_y), (x1, y1, x2, y2)))
|
||
|
||
# [删除] 之前生成的单张 _for_vl.jpg 逻辑已移除
|
||
|
||
# 保存标记图 (_flag.jpg)
|
||
if save_debug:
|
||
flag_out_path = os.path.join(output_dir, f"{stem}_flag{ext}")
|
||
save_image(flag_out_path, debug_img)
|
||
logger.info(f" Saved Debug Image: {flag_out_path}")
|
||
|
||
# 保存 _vl.jpg
|
||
vl_out_path = os.path.join(output_dir, f"{stem}_vl{ext}")
|
||
save_image(vl_out_path, vl_img)
|
||
logger.info(f" Step [2.2/VL] 已保存带有绿色方框的图片: {vl_out_path}")
|
||
|
||
# 保存 .json
|
||
import json
|
||
json_out_path = os.path.join(output_dir, f"{stem}.json")
|
||
with open(json_out_path, 'w', encoding='utf-8') as f:
|
||
json.dump(json_data, f, ensure_ascii=False, indent=4)
|
||
logger.info(f" Step [2.3/JSON] 已保存场站坐标元数据: {json_out_path}")
|
||
|
||
return results
|