252 lines
6.8 KiB
Python
252 lines
6.8 KiB
Python
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from mower.utils import typealias as tp
|
|
from mower.utils.log import logger
|
|
from mower.utils.path import get_path
|
|
|
|
|
|
def bytes2img(
|
|
img_data: bytes,
|
|
gray: bool = False,
|
|
bg: tuple[int, int, int, int] | str = "BLACK",
|
|
) -> tp.Image | tp.GrayImage:
|
|
"""解码图片
|
|
|
|
Args:
|
|
img_data (bytes): 图片的数据
|
|
gray (bool, optional): 返回灰图. Defaults to False.
|
|
bg (tuple[int, int, int, int] | str, optional): 读取透明PNG时设置的背景颜色. Defaults to "BLACK".
|
|
|
|
Returns:
|
|
tp.Image | tp.GrayImage: 图片
|
|
"""
|
|
img = cv2.imdecode(np.frombuffer(img_data, np.uint8), 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: # 彩色PNG
|
|
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: # 彩色JPG
|
|
if gray:
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
def img2bytes(
|
|
img: tp.Image | tp.GrayImage,
|
|
format: Literal["png", "jpg"] = "jpg",
|
|
quality: int = 75,
|
|
) -> bytes:
|
|
"""编码图片
|
|
|
|
Args:
|
|
img (tp.Image | tp.GrayImage): 图片
|
|
format (Literal["png", "jpg"], optional): 图片格式. Defaults to "jpg".
|
|
quality (int, optional): JPG质量. Defaults to 75.
|
|
|
|
Returns:
|
|
bytes: 图片数据
|
|
"""
|
|
if len(img.shape) == 3:
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
if format == "jpg":
|
|
data = cv2.imencode(".jpg", img, (cv2.IMWRITE_JPEG_QUALITY, quality))[1]
|
|
else:
|
|
data = cv2.imencode(".png", img)[1]
|
|
return data
|
|
|
|
|
|
def read_file(filename: str | Path) -> bytes:
|
|
"""读取图片数据，使用LRU Cache缓存
|
|
|
|
Args:
|
|
filename (str | Path): 文件名
|
|
|
|
Returns:
|
|
bytes: 图片数据
|
|
"""
|
|
filename = Path(filename)
|
|
try:
|
|
filename = filename.relative_to(get_path("@install"))
|
|
except ValueError:
|
|
pass
|
|
logger.debug(filename)
|
|
return np.fromfile(filename, dtype=np.uint8)
|
|
|
|
|
|
def loadimg(
|
|
filename: str, gray: bool = False, bg: tuple[int, int, int, int] | str = "BLACK"
|
|
) -> tp.Image | tp.GrayImage:
|
|
"""加载图片
|
|
|
|
Args:
|
|
filename (str): 文件名
|
|
gray (bool, optional): 返回灰图. Defaults to False.
|
|
bg (tuple[int, int, int, int] | str, optional): 透明PNG的背景颜色. Defaults to "BLACK".
|
|
|
|
Returns:
|
|
tp.Image | tp.GrayImage: 图片
|
|
"""
|
|
return bytes2img(read_file(filename), gray, bg)
|
|
|
|
|
|
def loadres(res: tp.Res, gray: bool = False) -> tp.Image | tp.GrayImage:
|
|
"""加载resources目录下的图片
|
|
|
|
Args:
|
|
res (tp.Res): 资源名
|
|
gray (bool, optional): 返回灰图. Defaults to False.
|
|
|
|
Returns:
|
|
tp.Image | tp.GrayImage: 图片
|
|
"""
|
|
res = f"@install/mower/resources/{res}"
|
|
if not res.endswith(".jpg"):
|
|
res += ".png"
|
|
return loadimg(get_path(res), gray)
|
|
|
|
|
|
loadres = lru_cache(maxsize=256)(loadres) # 使用装饰器没有类型提示
|
|
|
|
|
|
def load_static(name: str, gray: bool = False) -> tp.Image | tp.GrayImage:
|
|
"""加载static目录下的图片
|
|
|
|
Args:
|
|
name (str): 资源名
|
|
gray (bool, optional): 返回灰图. Defaults to False.
|
|
|
|
Returns:
|
|
tp.Image | tp.GrayImage: 图片
|
|
"""
|
|
name = f"@install/mower/static/{name}"
|
|
if not name.endswith(".jpg"):
|
|
name += ".png"
|
|
return loadimg(get_path(name), gray)
|
|
|
|
|
|
def load_static_dir(
|
|
subdir: str, gray: bool = True, str2int: bool = True
|
|
) -> dict[str, tp.Image | tp.GrayImage]:
|
|
"""加载static的一个子文件夹下的所有图片
|
|
|
|
Args:
|
|
subdir (str): 子文件夹
|
|
gray (bool, optional): 返回灰图. Defaults to False.
|
|
str2int (bool, optional): 字典键名转化为数字. Defaults to False.
|
|
|
|
Returns:
|
|
dict[str, tp.Image | tp.GrayImage]: 图片字典
|
|
"""
|
|
result = {}
|
|
for p in get_path(f"@install/mower/static/{subdir}").iterdir():
|
|
key = p.stem
|
|
if str2int:
|
|
try:
|
|
key = int(key)
|
|
except ValueError:
|
|
pass
|
|
result[key] = loadimg(p, gray)
|
|
return result
|
|
|
|
|
|
def thres2(img: tp.GrayImage, thresh: int) -> tp.GrayImage:
|
|
"""二值化
|
|
|
|
Args:
|
|
img (tp.GrayImage): 彩图
|
|
thresh (int): 阈值
|
|
|
|
Returns:
|
|
tp.GrayImage: 灰图
|
|
"""
|
|
return cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)[1]
|
|
|
|
|
|
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 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
|
|
|
|
logger.debug(f"{ca=} {cb=} {diff=}")
|
|
plt.imshow(board)
|
|
plt.show()
|
|
|
|
return diff <= thresh
|
|
|
|
|
|
def diff_ratio(
|
|
img1: tp.GrayImage,
|
|
img2: tp.GrayImage,
|
|
thresh: int = 0,
|
|
ratio: float = 0.05,
|
|
draw: bool = False,
|
|
) -> bool:
|
|
"""计算两张灰图之间不同的像素所占比例
|
|
|
|
Args:
|
|
img1 (tp.GrayImage): 一张灰图
|
|
img2 (tp.GrayImage): 另一张灰图
|
|
thresh (int, optional): 认为像素有差别的阈值. Defaults to 0.
|
|
ratio (float, optional): 判定有差别的比例阈值. Defaults to 0.05.
|
|
draw (bool, optional): 画出有差异的像素. Defaults to False.
|
|
|
|
Returns:
|
|
bool: 两张灰图是否有差异
|
|
"""
|
|
h, w = img1.shape
|
|
diff = cv2.absdiff(img1, img2)
|
|
thres = thres2(diff, thresh)
|
|
result = cv2.countNonZero(thres) / w / h
|
|
|
|
if draw:
|
|
from matplotlib import pyplot as plt
|
|
|
|
logger.debug(f"{result=}")
|
|
plt.imshow(diff, cmap="gray", vmin=0, vmax=255)
|
|
plt.show()
|
|
plt.imshow(thres, cmap="gray", vmin=0, vmax=255)
|
|
plt.show()
|
|
|
|
return result > ratio
|