167 lines
4.5 KiB
Python
167 lines
4.5 KiB
Python
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import Union
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from mower.utils import typealias as tp
|
|
from mower.utils.log import logger, save_screenshot
|
|
from mower.utils.path import get_path
|
|
|
|
|
|
def bytes2img(data: bytes, gray: bool = False) -> Union[tp.Image, tp.GrayImage]:
|
|
"""bytes -> image"""
|
|
if gray:
|
|
return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_GRAYSCALE)
|
|
else:
|
|
return cv2.cvtColor(
|
|
cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR),
|
|
cv2.COLOR_BGR2RGB,
|
|
)
|
|
|
|
|
|
def img2bytes(img: tp.Image) -> bytes:
|
|
"""image -> bytes"""
|
|
return cv2.imencode(
|
|
".jpg",
|
|
cv2.cvtColor(img, cv2.COLOR_RGB2BGR),
|
|
[int(cv2.IMWRITE_JPEG_QUALITY), 75],
|
|
)[1]
|
|
|
|
|
|
def res2path(res: tp.Res) -> Path:
|
|
res = f"@install/mower/resources/{res}"
|
|
if not res.endswith(".jpg"):
|
|
res += ".png"
|
|
return get_path(res)
|
|
|
|
|
|
def loadres(res: tp.Res, gray: bool = False) -> Union[tp.Image, tp.GrayImage]:
|
|
return loadimg(res2path(res), gray)
|
|
|
|
|
|
@lru_cache(maxsize=128)
|
|
def loadimg(
|
|
filename: str, gray: bool = False, bg: tuple[int] = (255, 255, 255, 255)
|
|
) -> Union[tp.Image, tp.GrayImage]:
|
|
"""load image from file"""
|
|
logger.debug(filename)
|
|
img_data = np.fromfile(filename, dtype=np.uint8)
|
|
img = cv2.imdecode(img_data, cv2.IMREAD_UNCHANGED)
|
|
|
|
if len(img.shape) == 2:
|
|
if gray:
|
|
return img
|
|
else:
|
|
return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
|
elif img.shape[2] == 4:
|
|
pim = Image.fromarray(img)
|
|
pbg = Image.new("RGBA", pim.size, bg)
|
|
pbg.paste(pim, (0, 0), pim)
|
|
if gray:
|
|
return np.array(pbg.convert("L"))
|
|
else:
|
|
return cv2.cvtColor(np.array(pbg.convert("RGB")), cv2.COLOR_BGR2RGB)
|
|
else:
|
|
if gray:
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
def thres2(img: tp.GrayImage, thresh: int) -> tp.GrayImage:
|
|
"""binarization of images"""
|
|
_, ret = cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)
|
|
return ret
|
|
|
|
|
|
# def thres0(img: tp.Image, thresh: int) -> tp.Image:
|
|
# """ delete pixel, filter: value > thresh """
|
|
# ret = img.copy()
|
|
# if len(ret.shape) == 3:
|
|
# # ret[rgb2gray(img) <= thresh] = 0
|
|
# z0 = ret[:, :, 0]
|
|
# z1 = ret[:, :, 1]
|
|
# z2 = ret[:, :, 2]
|
|
# _ = (z0 <= thresh) | (z1 <= thresh) | (z2 <= thresh)
|
|
# z0[_] = 0
|
|
# z1[_] = 0
|
|
# z2[_] = 0
|
|
# else:
|
|
# ret[ret <= thresh] = 0
|
|
# return ret
|
|
|
|
|
|
# def thres0(img: tp.Image, thresh: int) -> tp.Image: # not support multichannel image
|
|
# """ delete pixel which > thresh """
|
|
# _, ret = cv2.threshold(img, thresh, 255, cv2.THRESH_TOZERO)
|
|
# return ret
|
|
|
|
|
|
def rgb2gray(img: tp.Image) -> tp.GrayImage:
|
|
"""change image from rgb to gray"""
|
|
return cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
|
|
|
|
def scope2slice(scope: tp.Scope) -> tp.Slice:
|
|
"""((x0, y0), (x1, y1)) -> ((y0, y1), (x0, x1))"""
|
|
if scope is None:
|
|
return slice(None), slice(None)
|
|
return slice(scope[0][1], scope[1][1]), slice(scope[0][0], scope[1][0])
|
|
|
|
|
|
def cropimg(img: tp.Image, scope: tp.Scope) -> tp.Image:
|
|
"""crop image"""
|
|
return img[scope2slice(scope)]
|
|
|
|
|
|
def saveimg(img: tp.Image, folder):
|
|
del folder # 兼容2024.05旧版接口
|
|
save_screenshot(img2bytes(img))
|
|
|
|
|
|
def cmatch(
|
|
img1: tp.Image, img2: tp.Image, thresh: int = 10, draw: bool = False
|
|
) -> bool:
|
|
"比较平均色"
|
|
h, w, _ = img1.shape
|
|
ca = cv2.mean(img1)[:3]
|
|
cb = cv2.mean(img2)[:3]
|
|
diff = np.array(ca).astype(int) - np.array(cb).astype(int)
|
|
diff = np.max(np.maximum(diff, 0)) - np.min(np.minimum(diff, 0))
|
|
|
|
if draw:
|
|
board = np.zeros([h + 5, w * 2, 3], dtype=np.uint8)
|
|
board[:h, :w, :] = img1
|
|
board[h:, :w, :] = ca
|
|
board[:h, w:, :] = img2
|
|
board[h:, w:, :] = cb
|
|
|
|
from matplotlib import pyplot as plt
|
|
|
|
plt.imshow(board)
|
|
plt.show()
|
|
|
|
return diff <= thresh
|
|
|
|
|
|
def diff_ratio(
|
|
img1: tp.GrayImage,
|
|
img2: tp.GrayImage,
|
|
thresh: int = 0,
|
|
ratio: float = 0.05,
|
|
) -> bool:
|
|
"""计算两张灰图之间不同的像素所占比例
|
|
|
|
Args:
|
|
img1: 一张灰图
|
|
img2: 另一张灰图
|
|
thresh: 认为有差别的阈值
|
|
ratio: 有差别的像素比例阈值
|
|
"""
|
|
h, w = img1.shape
|
|
diff = cv2.absdiff(img1, img2)
|
|
diff = thres2(diff, thresh)
|
|
return cv2.countNonZero(diff) > ratio * w * h
|