status-desktop/test/e2e/scripts/tools/image.py

import logging
import time
import typing
from datetime import datetime

import allure
import cv2
import numpy as np
import pytesseract
from PIL import ImageGrab

import configs
import constants
import driver
from configs.system import IS_LIN
from scripts.tools.ocv import Ocv
from scripts.utils.system_path import SystemPath

LOG = logging.getLogger(__name__)


class Image:

    def __init__(self, object_name: dict):
        self.object_name = object_name
        self._view = None

    @property
    @allure.step('Get image view')
    def view(self) -> np.ndarray:
        return self._view

    @property
    @allure.step('Get image height')
    def height(self) -> int:
        return self.view.shape[0]

    @property
    @allure.step('Get image width')
    def width(self) -> int:
        return self.view.shape[1]

    @property
    @allure.step('Get image is grayscale')
    def is_grayscale(self) -> bool:
        return self.view.ndim == 2

    @allure.step('Set image in grayscale')
    def set_grayscale(self) -> 'Image':
        if not self.is_grayscale:
            self._view = cv2.cvtColor(self.view, cv2.COLOR_BGR2GRAY)
        return self

    @allure.step('Grab image view from object')
    def update_view(self):
        LOG.debug(f'Image view was grab from: {self.object_name}')
        rect = driver.object.globalBounds(driver.waitForObject(self.object_name))
        img = ImageGrab.grab(
            bbox=(rect.x, rect.y, rect.x + rect.width, rect.y + rect.height),
            xdisplay=configs.system.DISPLAY if IS_LIN else None
        )
        self._view = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)

    @allure.step('Save image')
    def save(self, path: SystemPath, force: bool = False):
        path.parent.mkdir(parents=True, exist_ok=True)
        if path.exists() and not force:
            raise FileExistsError(path)
        if self.view is None:
            self.update_view()
        cv2.imwrite(str(path), self.view)

    @allure.step('Compare images')
    def compare(
            self, expected: np.ndarray, threshold: float = 0.99) -> bool:
        if self.view is None:
            self.update_view()
        correlation = Ocv.compare_images(self.view, expected)
        result = correlation >= threshold
        LOG.info(f'Images equals on: {abs(round(correlation, 4) * 100)}%')

        if result:
            LOG.info(f'Screenshot comparison passed')
        else:
            configs.testpath.TEST_ARTIFACTS.mkdir(parents=True, exist_ok=True)
            diff = Ocv.draw_contours(self.view, expected)

            actual_fp = configs.testpath.TEST_ARTIFACTS / f'actual_image.png'
            expected_fp = configs.testpath.TEST_ARTIFACTS / f'expected_image.png'
            diff_fp = configs.testpath.TEST_ARTIFACTS / f'diff_image.png'

            self.save(actual_fp, force=True)
            cv2.imwrite(str(expected_fp), expected)
            cv2.imwrite(str(diff_fp), diff)

            allure.attach(name='actual', body=actual_fp.read_bytes(), attachment_type=allure.attachment_type.PNG)
            allure.attach(name='expected', body=expected_fp.read_bytes(), attachment_type=allure.attachment_type.PNG)
            allure.attach(name='diff', body=diff_fp.read_bytes(), attachment_type=allure.attachment_type.PNG)

            LOG.info(
                f"Screenshot comparison failed.\n"
                f"Actual, Diff and Expected screenshots are saved:\n"
                f"{configs.testpath.TEST_ARTIFACTS.relative_to(configs.testpath.ROOT)}.")
        return result

    @allure.step('Crop image')
    def crop(self, rect: driver.UiTypes.ScreenRectangle):
        assert rect.x + rect.width < self.width
        assert rect.y + rect.height < self.height
        self._view = self.view[rect.y: (rect.y + rect.height), rect.x: (rect.x + rect.width)]

    @allure.step('Parse text on image')
    def to_string(self, custom_config: str):
        text: str = pytesseract.image_to_string(self.view, config=custom_config)
        LOG.debug(f'Text on image: {text}')
        return text

    @allure.step('Verify: Image contains text: {1}')
    def has_text(self, text: str, criteria: str, crop: driver.UiTypes.ScreenRectangle = None) -> bool:
        self.update_view()
        if crop:
            self.crop(crop)

        # Search text on image converted in gray color
        self.set_grayscale()
        fp_gray = configs.testpath.TEST_ARTIFACTS / f'search_region_in_gray_color.png'
        self.save(fp_gray, force=True)
        if text.lower() in self.to_string(criteria).lower():
            allure.attach(name='search_region', body=fp_gray.read_bytes(), attachment_type=allure.attachment_type.PNG)
            return True

        # Search text on image with inverted color
        self._view = cv2.bitwise_not(self.view)
        fp_invert = configs.testpath.TEST_ARTIFACTS / f'search_region_in_inverted_color.png'
        self.save(fp_invert, force=True)
        if text.lower() in self.to_string(criteria).lower():
            allure.attach(name='search_region', body=fp_invert.read_bytes(), attachment_type=allure.attachment_type.PNG)
            return True
        return False

    @allure.step('Search color on image')
    def has_color(self, color: constants.Color, denoise: int = 10, crop: driver.UiTypes.ScreenRectangle = None) -> bool:
        self.update_view()
        if crop:
            self.crop(crop)

        initial_view = configs.testpath.TEST_ARTIFACTS / f'{color.name}.png'
        self.save(initial_view)
        allure.attach(name='search_region', body=initial_view.read_bytes(), attachment_type=allure.attachment_type.PNG)

        contours = self._get_color_contours(color, denoise, apply=True)

        mask_view = configs.testpath.TEST_ARTIFACTS / f'{color.name}_mask.png'
        self.save(mask_view)
        allure.attach(name='contours', body=mask_view.read_bytes(), attachment_type=allure.attachment_type.PNG)

        self._view = None
        return len(contours) >= 1

    @allure.step('Apply contours with found color on image')
    def _get_color_contours(
            self,
            color: constants.Color,
            denoise: int = 10,
            apply: bool = False
    ) -> typing.List[driver.UiTypes.ScreenRectangle]:
        if not self.is_grayscale:
            view = cv2.cvtColor(self.view, cv2.COLOR_BGR2HSV)
        else:
            view = self.view
        boundaries = constants.boundaries[color]

        if color == constants.Color.RED:
            mask = None
            for bond in boundaries:
                lower_range = np.array(bond[0])
                upper_range = np.array(bond[1])
                _mask = cv2.inRange(view, lower_range, upper_range)
                mask = _mask if mask is None else mask + _mask
        else:
            lower_range = np.array(boundaries[0])
            upper_range = np.array(boundaries[1])
            mask = cv2.inRange(view, lower_range, upper_range)

        contours = []
        _contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        _contours = _contours[0] if len(_contours) == 2 else _contours[1]
        for _contour in _contours:
            x, y, w, h = cv2.boundingRect(_contour)
            # To remove small artifacts, less than denoise variable value
            if w * h < denoise:
                continue
            contours.append(driver.UiTypes.ScreenRectangle(x, y, w, h))

        if apply:
            self._view = cv2.bitwise_and(self.view, self.view, mask=mask)
            for contour in contours:
                cv2.rectangle(
                    self.view,
                    (contour.x, contour.y),
                    (contour.x + contour.width, contour.y + contour.height),
                    (36, 255, 12), 2)

        return contours


@allure.step('Compare images')
def compare(actual: Image,
            expected: typing.Union[str, SystemPath, Image],
            threshold: float = 0.99,
            timout_sec: int = 1
            ):
    expected_fp = None
    if isinstance(expected, str):
        expected_fp = configs.testpath.TEST_VP / configs.system.OS_ID / expected
        if not expected_fp.exists():
            expected_fp = configs.testpath.TEST_VP / expected
        expected = expected_fp
    if isinstance(expected, SystemPath):
        assert expected.exists(), f'File: {expected} not found'
        expected = cv2.imread(str(expected))
    else:
        expected = expected.view
    start = datetime.now()
    while not actual.compare(expected, threshold):
        time.sleep(1)
        if (datetime.now() - start).seconds > timout_sec:
            if configs.UPDATE_VP_ON_FAIL and expected_fp is not None:
                actual.save(expected_fp, force=True)
                LOG.warning(f'VP file updated: {expected_fp}')
                break
            else:
                raise AssertionError('Images comparison failed')
    LOG.info(f'Screenshot comparison passed')