zhang_crawl_bio/utils/chrome_driver/cv_common.py

import asyncio
import time
import cv2
from cv import detect
from DrissionPage import ChromiumPage
from DrissionPage.common import Actions,Keys

from sqlmodel import SQLModel, Field, Column, JSON
from sqlalchemy.types import LargeBinary
from typing import List, Optional,Coroutine

class ImageMatchResult(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    screen_img_data: bytes = Field(sa_column=Column(LargeBinary))
    template_img_path: str
    template_img_h: int
    template_img_w: int
    threshold: float
    match_min_val: float
    match_max_val: float
    min_location: list = Field(sa_column=Column(JSON))
    max_location: list = Field(sa_column=Column(JSON))
    max_location_center: List[int] = Field(sa_column=Column(JSON))

class CVPage:
    def __init__(self, tab:ChromiumPage=None,) -> None:
        self.tab =  tab

    @classmethod
    def match_img_in_screen(cls, screen, target_path) -> ImageMatchResult:
        screen_img = detect.read_image(screen)
        template_img = detect.read_image(target_path)

        # 获取模板图像的宽度和高度
        w, h = template_img.shape[1], template_img.shape[0]

        # 模板匹配
        result = cv2.matchTemplate(screen_img, template_img, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_location, max_location = cv2.minMaxLoc(result)
        match_result = ImageMatchResult(
            screen_img_path=screen_img,
            template_img_path=target_path,
            template_img_h=h,
            template_img_w=w,
            match_min_val=min_val,
            match_max_val=max_val,
            min_location=min_location,
            max_location=max_location,
            max_location_center= [max_location[0] + w // 2, max_location[1] + h // 2]
        )
        return match_result
    
    @classmethod
    def match_img_in_screen_region(cls, screen, target_path, region=None) -> ImageMatchResult:
        screen_img: cv2.typing.MatLike = detect.read_image(screen)
        template_img: cv2.typing.MatLike = detect.read_image(target_path)

        # 获取屏幕图像的尺寸
        screen_h, screen_w = screen_img.shape[:2]

        # 如果指定了region，计算实际的匹配区域
        if region:
            x1, y1, x2, y2 = region
            x1, y1 = int(x1 * screen_w), int(y1 * screen_h)
            x2, y2 = int(x2 * screen_w), int(y2 * screen_h)
            roi = screen_img[y1:y2, x1:x2]
        else:
            roi = screen_img

        # 获取模板图像的宽度和高度
        w, h = template_img.shape[1], template_img.shape[0]
        # cv2.imwrite('test1.png', roi)
        # 模板匹配
        result = cv2.matchTemplate(roi, template_img, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_location, max_location = cv2.minMaxLoc(result)

        match_result = ImageMatchResult(
            screen_img_path=screen if type(screen) == str else None,
            template_img_path=target_path,
            template_img_h=h,
            template_img_w=w,
            match_min_val=min_val,
            match_max_val=max_val,
            min_location=min_location,
            max_location=max_location,
            max_location_center=[max_location[0] + w // 2, max_location[1] + h // 2]
        )
        return match_result
    
    def show_match_result(screen, match_result:ImageMatchResult, threshold=0.8, imshow=True):
        # 读取源图像和目标图像
        source_img: cv2.typing.MatLike = detect.read_image(screen)
        
        w = match_result.template_img_w   
        h = match_result.template_img_h     
        # 获取匹配结果
        max_val = match_result.match_max_val
        max_loc = match_result.max_location
        
        if max_val >= threshold:  # 如果最大相似度大于或等于阈值
            top_left = max_loc
            bottom_right = (top_left[0] + w, top_left[1] + h)
            label = f"Similarity: {max_val:.2f}"
            cv2.putText(source_img, label, (top_left[0], top_left[1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            # 绘制矩形框
            cv2.rectangle(source_img, top_left, bottom_right, (0, 255, 0), 2)
        #     print(f"匹配位置: {top_left,bottom_right}, 相似度: {max_val}")

        # else:
        #     print("没有找到相似度足够高的匹配.")
            
        if imshow:
            # 创建一个窗口来显示结果
            cv2.namedWindow('Detected', cv2.WINDOW_NORMAL)
            cv2.imshow('Detected', source_img)
            
            # 等待按键
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        else:
            return source_img
    
    def wait_for_image_match(self, target_path, threshold=0.9, timeout=30, interval=0.3, raise_error=False):  
        """  
        等待直到图片匹配度超过给定阈值或超时。  
  
        :param target_path: 目标图片的路径  
        :param threshold: 匹配度阈值，默认为0.8  
        :param timeout: 等待超时时间（秒），默认为30秒  
        :param interval: 每次检查之间的间隔时间（秒），默认为0.5秒  
        :return: 如果找到匹配的图片，则返回ImageMatchResult，否则返回None  
        """  
        end_time = time.time() + timeout  
        while time.time() < end_time:  
            match_result = self.match_img_in_screen(target_path)  
            if match_result.match_max_val >= threshold:  
                return match_result  
            time.sleep(interval)  
        if raise_error:  
            raise Exception(f'超时未找到匹配的图片 {target_path} ，阈值为{threshold}')
        return None  # 超时未找到符合条件的匹配 

    async def wait_for_image_match_async(self, screen, target_path, threshold=0.9, timeout=30, interval=0.3, raise_error=False)->ImageMatchResult:
            """
            异步等待直到图片匹配度超过给定阈值或超时。
            """
            end_time = time.time() + timeout
            while time.time() < end_time:
                if isinstance(screen, Coroutine):
                    print(screen)
                    screen = await screen()
                match_result = self.match_img_in_screen(screen, target_path)
                if match_result.match_max_val >= threshold:
                    return match_result
                await asyncio.sleep(interval)
            if raise_error:
                raise Exception(f'超时未找到匹配的图片 {target_path} ，阈值为{threshold}')
            return None  # 超时未找到符合条件的匹配    
    def wait_for_image_disappear(self, target_path, threshold=0.9, timeout=30, interval=0.3, raise_error=False):  
        """  
        等待直到目标图片从屏幕上消失或超时。  
  
        :param target_path: 目标图片的路径  
        :param threshold: 匹配度阈值，当匹配度低于此值时认为图片已消失，默认为0.3  
        :param timeout: 等待超时时间（秒），默认为30秒  
        :param interval: 每次检查之间的间隔时间（秒），默认为0.5秒  
        :return: 如果图片消失，则返回True，否则返回False  
        """  
        end_time = time.time() + timeout  
        while time.time() < end_time:  
            match_result = self.match_img_in_screen(target_path)  
            if match_result.match_max_val < threshold:  
                return True  # 图片已消失  
            time.sleep(interval)  
        if raise_error:  
            raise Exception(f'超时未找到匹配的图片 {target_path} ，阈值为{threshold}')
        return None  # 超时未找到符合条件的匹配 

    async def wait_for_image_disappear_async(self, screen, target_path, threshold=0.9, timeout=30, interval=0.3, raise_error=False):
            """
            异步等待直到目标图片从屏幕上消失或超时。
            """
            end_time = time.time() + timeout
            while time.time() < end_time:
                match_result = self.match_img_in_screen(screen, target_path)
                if match_result.match_max_val < threshold:
                    return True  # 图片已消失
                await asyncio.sleep(interval)
            if raise_error:
                raise Exception(f'超时未找到匹配的图片 {target_path} ，阈值为{threshold}')
            return False  # 超时未找到符合条件的匹配