157 lines
4.5 KiB
Python
157 lines
4.5 KiB
Python
# coding=utf-8
|
||
import os
|
||
import time
|
||
import logging
|
||
|
||
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, path=None):
|
||
"""
|
||
获取屏幕截图并保存
|
||
"""
|
||
if not path:
|
||
from Config.Config import TEMP_IMAGE_DIR
|
||
path = TEMP_IMAGE_DIR
|
||
|
||
if not os.path.exists(path):
|
||
os.makedirs(path)
|
||
|
||
full_path = os.path.join(path, filename)
|
||
d.screenshot(full_path)
|
||
return full_path
|
||
|
||
import hashlib
|
||
import numpy as np
|
||
import cv2
|
||
|
||
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
|
||
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:
|
||
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 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 get_image_md5(image_path):
|
||
"""
|
||
计算图片的 MD5 值
|
||
"""
|
||
import hashlib
|
||
if not os.path.exists(image_path):
|
||
return ""
|
||
with open(image_path, 'rb') as f:
|
||
data = f.read()
|
||
return hashlib.md5(data).hexdigest()
|
||
|
||
def draw_rectangles(image_path, points, output_path=None):
|
||
"""
|
||
使用 OpenCV 在图片上绘制矩形框 (计算机图形学绘制)
|
||
:param image_path: 原始图片路径
|
||
:param points: 坐标点列表，格式如 [[x1, y1, x2, y2], ...] 或 [(x, y), ...]
|
||
:param output_path: 输出图片路径，如果不指定则覆盖原图
|
||
"""
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
from Config.Setting import DRAW_DEBUG_BOXES, DEBUG_BOX_COLOR, DEBUG_BOX_THICKNESS
|
||
|
||
if not DRAW_DEBUG_BOXES:
|
||
return image_path
|
||
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
return image_path
|
||
|
||
for p in points:
|
||
if len(p) == 4: # 矩形 [x1, y1, x2, y2]
|
||
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: # 点 (x, y)，画一个小圆圈或正方形表示点击位
|
||
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
|
||
cv2.imwrite(save_path, img)
|
||
return save_path
|
||
except Exception as e:
|
||
logger.error(f"绘制矩形框失败: {e}")
|
||
return image_path
|