'commit'
This commit is contained in:
HuangHai
324
Apps/AiTeJiYiChong/Kit.py
Normal file
324
Apps/AiTeJiYiChong/Kit.py
Normal file
@@ -0,0 +1,324 @@
|
||||
import logging
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import time
|
||||
import json
|
||||
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 read_image(path):
|
||||
"""读取图片，支持中文路径"""
|
||||
try:
|
||||
return cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
|
||||
except Exception as e:
|
||||
logger.info(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"Found template at ({center_x}, {center_y}) with confidence {max_val:.2f}")
|
||||
d.click(center_x, center_y)
|
||||
return True
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
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)],
|
||||
"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
|
||||
Reference in New Issue
Block a user