Source code for airtest.aircv.utils

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
import time
import numpy as np
from PIL import Image

from airtest.utils.logger import get_logger
from .error import TemplateInputError

LOGGING = get_logger(__name__)





[docs]def generate_result(middle_point, pypts, confi): """Format the result: 定义图像识别结果格式.""" ret = dict(result=middle_point, rectangle=pypts, confidence=confi) return ret
[docs]def check_image_valid(im_source, im_search): """Check if the input images valid or not.""" if im_source is not None and im_source.any() and im_search is not None and im_search.any(): return True else: return False
[docs]def img_mat_rgb_2_gray(img_mat): """ Turn img_mat into gray_scale, so that template match can figure the img data. "print(type(im_search[0][0])") can check the pixel type. """ assert isinstance(img_mat[0][0], np.ndarray), "input must be instance of np.ndarray" return cv2.cvtColor(img_mat, cv2.COLOR_BGR2GRAY)
[docs]def img_2_string(img): _, png = cv2.imencode('.png', img) return png.tostring()
[docs]def string_2_img(pngstr): nparr = np.frombuffer(pngstr, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) return img
[docs]def pil_2_cv2(pil_image): open_cv_image = np.array(pil_image) # Convert RGB to BGR (method-1): open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR) # Convert RGB to BGR (method-2): # b, g, r = cv2.split(open_cv_image) # open_cv_image = cv2.merge([r, g, b]) return open_cv_image
[docs]def cv2_2_pil(cv2_image): cv2_im = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) pil_im = Image.fromarray(cv2_im) return pil_im
[docs]def compress_image(pil_img, path, quality, max_size=None): """ Save the picture and compress :param pil_img: PIL image :param path: save path :param quality: the image quality, integer in range [1, 99] :param max_size: the maximum size of the picture, e.g 1200 :return: """ if max_size: # The picture will be saved in a size <= max_size*max_size pil_img.thumbnail((max_size, max_size), Image.LANCZOS) quality = int(round(quality)) if quality <= 0 or quality >= 100: raise Exception("SNAPSHOT_QUALITY (" + str(quality) + ") should be an integer in the range [1,99]") pil_img.save(path, quality=quality, optimize=True)