185 lines
6.9 KiB
Python
185 lines
6.9 KiB
Python
import asyncio
|
||
import io,base64
|
||
import logging,time
|
||
from typing import Optional, Tuple
|
||
from PIL import Image, ImageChops
|
||
import numpy as np
|
||
from pydantic import BaseModel
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class VisualChangeDetector:
|
||
"""视觉变化检测器（基于屏幕截图比对）"""
|
||
|
||
def __init__(
|
||
self,
|
||
similarity_threshold: float = 0.95,
|
||
pixel_change_threshold: int = 1000,
|
||
resize_width: Optional[int] = 800
|
||
):
|
||
"""
|
||
Args:
|
||
similarity_threshold: 相似度阈值（0-1）
|
||
pixel_change_threshold: 变化像素数量阈值
|
||
resize_width: 缩放宽度（提升性能）
|
||
"""
|
||
self.similarity_threshold = similarity_threshold
|
||
self.pixel_threshold = pixel_change_threshold
|
||
self.resize_width = resize_width
|
||
|
||
async def capture_screenshot(self, page, upload_service=None,logger=None) -> Image.Image:
|
||
"""捕获页面截图（自动优化）"""
|
||
screenshot_bytes = await page.screenshot(timeout=10000,type='jpeg',quality=50,full_page=False,animations='disabled')
|
||
img = Image.open(io.BytesIO(screenshot_bytes))
|
||
|
||
# 统一缩放尺寸（提升比对性能）
|
||
if self.resize_width:
|
||
w, h = img.size
|
||
new_h = int(h * (self.resize_width / w))
|
||
img = img.resize((self.resize_width, new_h))
|
||
# if upload_service:
|
||
# img_byte_arr = io.BytesIO()
|
||
# img.save(img_byte_arr, format='JPEG')
|
||
# binary_data = img_byte_arr.getvalue()
|
||
# if not binary_data:
|
||
# return
|
||
# file_name = f"screenshot_{int(time.time())}.jpg"
|
||
# base64_str = base64.b64encode(binary_data).decode('ascii')
|
||
# if logger:
|
||
# logger.info(f'upload screenshot to {file_name}')
|
||
# else:
|
||
# logging.info(f'upload screenshot to {file_name}')
|
||
# try:
|
||
# await upload_service.upload_file('',file_name,base64_content=base64_str)
|
||
# except Exception as e:
|
||
# if logger:
|
||
# logger.error(f"Failed to upload screenshot to {file_name}: {e}")
|
||
# else:
|
||
# logging.error(f"Failed to upload screenshot to {file_name}: {e}")
|
||
# return img.convert('RGB')
|
||
return img.convert('RGB') # 确保RGB模式
|
||
|
||
def calculate_change(
|
||
self,
|
||
img1: Image.Image,
|
||
img2: Image.Image
|
||
) -> Tuple[float, Image.Image]:
|
||
"""
|
||
计算两图差异
|
||
|
||
Returns:
|
||
tuple: (相似度百分比, 差异图)
|
||
"""
|
||
# 确保尺寸一致
|
||
if img1.size != img2.size:
|
||
img2 = img2.resize(img1.size)
|
||
|
||
arr1 = np.array(img1) # shape: (height, width, 3)
|
||
arr2 = np.array(img2) # shape: (height, width, 3)
|
||
|
||
# 计算绝对差异（每个通道单独计算）
|
||
diff = np.abs(arr1.astype(int) - arr2.astype(int))
|
||
|
||
# 变化像素计算（任一通道有变化即视为该像素变化）
|
||
# 先检查每个像素的3个通道是否有变化
|
||
pixel_changed = np.any(diff > 0, axis=2) # shape: (height, width)
|
||
changed_pixels = np.sum(pixel_changed) # 变化像素总数
|
||
|
||
# 总像素数（不需要除以3，因为pixel_changed已经是像素级别的判断）
|
||
total_pixels = arr1.shape[0] * arr1.shape[1]
|
||
|
||
# 生成差异图（可视化用）
|
||
diff_img = ImageChops.difference(img1, img2)
|
||
|
||
# 计算相似度
|
||
similarity = 1 - (changed_pixels / total_pixels)
|
||
|
||
return similarity, diff_img
|
||
|
||
async def detect_change(
|
||
self,
|
||
browser_session,
|
||
reference_img: Optional[Image.Image] = None,
|
||
max_attempts: int = 5,
|
||
attempt_interval: float = 1.5,
|
||
upload_service = None,
|
||
logger = None
|
||
) -> Tuple[bool, Optional[Image.Image], Optional[Image.Image]]:
|
||
"""
|
||
检测视觉变化
|
||
|
||
Args:
|
||
page: 浏览器页面对象
|
||
reference_img: 基准截图（None则自动捕获）
|
||
max_attempts: 最大检测次数
|
||
attempt_interval: 检测间隔（秒）
|
||
|
||
Returns:
|
||
tuple: (是否变化, 基准截图, 差异图)
|
||
"""
|
||
# 首次捕获基准图
|
||
if reference_img is None:
|
||
page = await browser_session.get_current_page()
|
||
reference_img = await self.capture_screenshot(page,upload_service=upload_service,logger=logger)
|
||
logger.info("start detect change")
|
||
for attempt in range(max_attempts):
|
||
await asyncio.sleep(attempt_interval)
|
||
|
||
# 捕获当前截图
|
||
page = await browser_session.get_current_page()
|
||
current_img = await self.capture_screenshot(page,upload_service=upload_service,logger=logger)
|
||
if not current_img:
|
||
logger.error(f"Failed to capture screenshot on attempt {attempt + 1}")
|
||
continue
|
||
# 计算变化
|
||
similarity, diff_img = self.calculate_change(reference_img, current_img)
|
||
logger.info(f"Attempt {attempt + 1}: 相似度 {similarity:.2f}, 变化像素 {np.sum(np.array(diff_img) > 0)}")
|
||
# 判断是否显著变化
|
||
if similarity < self.similarity_threshold:
|
||
diff_pixels = np.sum(np.array(diff_img) > 0)
|
||
if diff_pixels > self.pixel_threshold:
|
||
logger.info(f"视觉变化 detected (相似度: {similarity:.2f}, 变化像素: {diff_pixels})")
|
||
return True, reference_img, diff_img
|
||
|
||
return False, reference_img, None
|
||
|
||
class WaitForVisualChangeAction(BaseModel):
|
||
"""等待视觉变化的参数模型"""
|
||
timeout: int = 30
|
||
check_interval: float = 2
|
||
similarity_threshold: float = 0.95
|
||
pixel_threshold: int = 5000
|
||
|
||
async def wait_for_visual_change(
|
||
params: WaitForVisualChangeAction,
|
||
browser_session,
|
||
initial_screenshot: Optional[Image.Image] = None,
|
||
upload_service = None,
|
||
logger = None
|
||
) -> Tuple[bool, Optional[Image.Image], Optional[Image.Image]]:
|
||
"""
|
||
等待页面视觉变化（完整工作流）
|
||
|
||
Args:
|
||
params: 配置参数
|
||
browser_session: 浏览器会话
|
||
initial_screenshot: 初始截图（可选）
|
||
|
||
Returns:
|
||
tuple: (是否变化, 初始截图, 差异图)
|
||
"""
|
||
detector = VisualChangeDetector(
|
||
similarity_threshold=params.similarity_threshold,
|
||
pixel_change_threshold=params.pixel_threshold
|
||
)
|
||
|
||
max_attempts = int(params.timeout / params.check_interval)
|
||
|
||
return await detector.detect_change(
|
||
browser_session=browser_session,
|
||
reference_img=initial_screenshot,
|
||
max_attempts=max_attempts,
|
||
attempt_interval=params.check_interval,
|
||
upload_service=upload_service,
|
||
logger=logger,
|
||
) |