mirror of
https://github.com/MaaAssistantArknights/MaaAssistantArknights.git
synced 2026-07-15 17:30:27 +08:00
d025969c82
* fix: 修复 ImageCropper 循环异常退出的问题 * feat: 允许 ADB 为空 * fix: 修复 input 为空时的异常 * chore: Auto update by pre-commit hooks [skip changelog] * chore: 规范变量命名 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
471 lines
15 KiB
Python
471 lines
15 KiB
Python
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import colormatcher
|
||
import cv2
|
||
import numpy as np
|
||
from maa.controller import AdbController, Controller
|
||
from maa.define import MaaAdbScreencapMethodEnum
|
||
from maa.toolkit import Toolkit
|
||
from roimage import Roi, Roimage
|
||
|
||
# 初始化设备参数
|
||
# device_serial = "127.0.0.1:16384"
|
||
device_serial = None
|
||
adb_screencap_method = MaaAdbScreencapMethodEnum.Default
|
||
|
||
# 初始窗口大小 (width, height)
|
||
# window_size = (720, 1280) # 竖屏
|
||
window_size = (1280, 720) # 横屏
|
||
|
||
# 右键拖曳
|
||
move_start = (0, 0)
|
||
move_start_roi = (0, 0)
|
||
|
||
|
||
# 初始化 Roi
|
||
std_roimage: Optional[Roimage] = Roimage(window_size[0], window_size[1]) # 标准化截图
|
||
win_roimage: Roimage = Roimage(
|
||
0, 0, 0, 0, std_roimage
|
||
) # 相对 std_roimage ，窗口显示的区域
|
||
crop_list: list[Roimage] = [] # 相对 std_roimage ，需要裁剪的区域
|
||
|
||
# 左键裁剪ROI区域
|
||
crop_start: Optional[Roimage] = None # 相对 win_roimage ，正在裁剪的区域
|
||
crop_end = Roimage(0, 0, 0, 0, win_roimage).getRoiInRoot()
|
||
|
||
# 初始化参数
|
||
win_name = "image"
|
||
trackbars_name = "trackbars"
|
||
file_name = "image"
|
||
src_path = Path("./src")
|
||
files = [f.name for f in src_path.glob("*.png")] if src_path.exists() else []
|
||
|
||
# 显示 Roi
|
||
trackbars_img = np.ones((100, 400, 3), dtype=np.uint8) * 255
|
||
|
||
# -----------------------------------------------
|
||
|
||
|
||
# 截图参数
|
||
def set_screenshot_target_side(c: Controller):
|
||
# c.set_screenshot_target_long_side(1280)
|
||
c.set_screenshot_target_short_side(720)
|
||
|
||
|
||
# ROI 放大方法
|
||
def amplify(rect: list[int]) -> list[int]:
|
||
x, y, w, h = rect
|
||
return [x - 50, y - 50, w + 100, h + 100]
|
||
|
||
|
||
# 颜色匹配方法
|
||
def match_color(image) -> tuple[int, int, list[tuple[list[int]]]]:
|
||
# framework ColorMatch common method:
|
||
# 4.COLOR_BGR2RGB 40.COLOR_BGR2HSV 6.COLOR_BGR2GRAY
|
||
# colormatcher method:
|
||
# 1.Simple 2.RGBDistance
|
||
method = cv2.COLOR_BGR2RGB
|
||
reverse = cv2.COLOR_RGB2BGR # 不需要逆转时为None，用于显示 'MainColors' 窗口
|
||
cluster_colors = colormatcher.kmeansClusterColors(image, method)
|
||
return method, reverse, colormatcher.RGBDistance(cluster_colors)
|
||
|
||
|
||
# 获取 ADB 设备列表
|
||
def get_adb_devices_list() -> Optional[Controller]:
|
||
print("MaaToolkit search in progress...\n")
|
||
device_list = Toolkit.find_adb_devices()
|
||
if len(device_list):
|
||
for i, d in enumerate(device_list):
|
||
print(f"{i:>3} | {d.address:>21} | {d.name}")
|
||
|
||
if device_input := input("Please select the device (ENTER to pass): "):
|
||
i = int(device_input)
|
||
if 0 <= i < len(device_list):
|
||
device_serial = device_list[i].address
|
||
return AdbController(
|
||
adb_path=device_list[i].adb_path,
|
||
address=device_serial,
|
||
screencap_methods=adb_screencap_method,
|
||
)
|
||
|
||
|
||
# OpenCV 鼠标回调
|
||
# -events 鼠标事件（如按下鼠标左键，释放鼠标左键，鼠标移动等）
|
||
# -x x坐标
|
||
# -y y坐标
|
||
# -flags params 其他参数
|
||
def mouse(event, x, y, flags, param) -> None:
|
||
global crop_end
|
||
crop_end = Roimage(0, 0, x, y, win_roimage).getRoiInRoot()
|
||
crop(event, x, y, flags, param)
|
||
zoom(event, x, y, flags, param)
|
||
move(event, x, y, flags, param)
|
||
show_roi(crop_end)
|
||
|
||
|
||
def show_roi(roi: Roi):
|
||
trackbars_img.fill(255)
|
||
cv2.putText(
|
||
trackbars_img,
|
||
f"{roi.rectangle}",
|
||
(0, 60),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
1,
|
||
(0, 0, 0),
|
||
2,
|
||
)
|
||
cv2.imshow(trackbars_name, trackbars_img)
|
||
|
||
|
||
# 计算绘图四边形坐标
|
||
# -rectPts 相对于 std_roimage 的两点坐标 ((left,top),(right,bottom))
|
||
def count_draw_coordinate(rect_pts):
|
||
z = win_roimage.zoom
|
||
x, y = win_roimage.point
|
||
return (
|
||
(int(rect_pts[0][0] * z - x), int(rect_pts[0][1] * z - y)),
|
||
(int(rect_pts[1][0] * z - x), int(rect_pts[1][1] * z - y)),
|
||
)
|
||
|
||
|
||
# 绘图
|
||
# -rois 相对于 std_roimage 的 crop_roimage 列表
|
||
def draw(rois: list[Roimage]) -> None:
|
||
img = win_roimage.image.copy()
|
||
for roi in rois:
|
||
pt1, pt2 = count_draw_coordinate(roi.rectanglePoints)
|
||
cv2.rectangle(img, pt1, pt2, (0, 255, 0), 2)
|
||
cv2.imshow(win_name, img)
|
||
|
||
|
||
def crop(event, x, y, flags, param) -> None:
|
||
global crop_start, crop_end
|
||
if event == cv2.EVENT_LBUTTONDOWN: # 按下左键
|
||
# 记录（x,y）坐标
|
||
crop_start = Roimage(0, 0, x, y, win_roimage)
|
||
elif crop_start is not None and (
|
||
event == cv2.EVENT_LBUTTONUP
|
||
or (event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_LBUTTON))
|
||
): # 释放左键 或 按住左键拖曳
|
||
# 修正（x,y）坐标
|
||
w, h = win_roimage.size
|
||
x = max(0, min(w - 1, x))
|
||
y = max(0, min(h - 1, y))
|
||
# 记录（x,y）坐标
|
||
crop_end = crop_start.getCropRoi(x, y).getRoiInRoot()
|
||
# 绘图
|
||
rois: list[Roimage] = crop_list.copy()
|
||
rois.append(crop_end)
|
||
draw(rois)
|
||
# 保存 crop Roi
|
||
if event == cv2.EVENT_LBUTTONUP:
|
||
crop_list.append(crop_end.copy(std_roimage))
|
||
crop_start = None
|
||
|
||
|
||
# 计算缩放倍数
|
||
# -flag 鼠标滚轮上移或下移的标识
|
||
# -zoom 缩放倍数
|
||
# -step 缩放系数，滚轮每步缩放0.1
|
||
def count_zoom(flag, zoom: float, step: float = 0.1):
|
||
if flag > 0: # 滚轮上移
|
||
zoom += step
|
||
if zoom > 3: # 最多只能放大到 g_image 三倍大
|
||
zoom = 3
|
||
else: # 滚轮下移
|
||
zoom -= step
|
||
if zoom < 1: # 最多只能缩小到 g_image 大小
|
||
zoom = 1
|
||
return zoom
|
||
|
||
|
||
# 滚轮放大
|
||
def zoom(event, x, y, flags, param) -> None:
|
||
global win_roimage
|
||
if event == cv2.EVENT_MOUSEWHEEL: # 滚轮
|
||
z = count_zoom(flags, win_roimage.zoom)
|
||
wx, wy, w, h = win_roimage.rectangle
|
||
x = int((wx + x) * z / win_roimage.zoom - x)
|
||
y = int((wy + y) * z / win_roimage.zoom - y)
|
||
win_roimage = Roimage(w, h, x, y, win_roimage.parent, z)
|
||
draw(crop_list)
|
||
|
||
|
||
# 计算移动后的坐标
|
||
# -pt0 win_roimage 的原始坐标
|
||
# -pt1 鼠标按下右键时的坐标
|
||
# -pt2 鼠标当前坐标
|
||
def count_move_coordinate(pt0, pt1, pt2):
|
||
return pt0[0] + pt1[0] - pt2[0], pt0[1] + pt1[1] - pt2[1]
|
||
|
||
|
||
def move(event, x, y, flags, param) -> None:
|
||
global move_start, move_start_roi, win_roimage
|
||
if event == cv2.EVENT_RBUTTONDOWN:
|
||
move_start = (x, y)
|
||
move_start_roi = win_roimage.point
|
||
elif event == cv2.EVENT_RBUTTONUP or (
|
||
event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_RBUTTON)
|
||
):
|
||
x, y = count_move_coordinate(move_start_roi, move_start, (x, y))
|
||
win_roimage = Roimage(
|
||
win_roimage.width, win_roimage.height, x, y, std_roimage, win_roimage.zoom
|
||
)
|
||
draw(crop_list)
|
||
|
||
|
||
# 轨迹条回调
|
||
# -pos 轨迹条位置
|
||
def trackbar_change(pos) -> None:
|
||
pos = pos / 100 # get a scaling factor from trackbar pos
|
||
w = int(std_roimage.width * pos) # scale w
|
||
h = int(std_roimage.height * pos) # scale h
|
||
cv2.resizeWindow(win_name, w, h) # resize a window
|
||
|
||
|
||
def readfile(file: str):
|
||
file_path = src_path / file
|
||
print("src:", file_path.absolute())
|
||
return cv2.imread(str(file_path))
|
||
|
||
|
||
# 截图
|
||
def screenshot() -> Optional[np.ndarray]:
|
||
if not controller:
|
||
return None
|
||
print("Screenshot in progress...")
|
||
controller.post_screencap().wait()
|
||
image = controller.cached_image
|
||
print("Screenshot completed.")
|
||
return image
|
||
|
||
|
||
# 获取标准化的 Roimage
|
||
def get_std_roimage() -> Optional[Roimage]:
|
||
global file_name
|
||
image: Optional[np.ndarray] = None
|
||
if len(files):
|
||
file_name = files.pop(0)
|
||
image = readfile(file_name)
|
||
file_name = file_name.split(".")[0]
|
||
elif controller:
|
||
image = screenshot()
|
||
file_name = datetime.now().strftime("%H%M%S") # '%Y%m%d%H%M%S'
|
||
if image is None:
|
||
return None
|
||
height, width, _ = image.shape
|
||
roimage = Roimage(width, height)
|
||
roimage.image = image
|
||
return roimage
|
||
|
||
|
||
# 获取放大后的 Roi 四边形
|
||
# -roi: 需要放大的 Roi
|
||
def get_amplified_roi_rectangle(roi: Roimage) -> list[int]:
|
||
x, y, w, h = amplify(roi.rectangle)
|
||
return Roimage(w, h, x, y, roi.parent).rectangle
|
||
|
||
|
||
def is_window_visible(win_name: str, destroyed: bool = False) -> bool:
|
||
if destroyed:
|
||
return False
|
||
try:
|
||
return cv2.getWindowProperty(win_name, cv2.WND_PROP_VISIBLE) > 0
|
||
except cv2.error:
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# Print Help Infos
|
||
print(
|
||
"Usage: python3 main.py [device serial]\n"
|
||
"Put the images under ./src, and run this script, it will be auto converted to target size.\n"
|
||
"Hold down the left mouse button, drag mouse to select a ROI.\n"
|
||
"Hold down the right mouse button, drag mouse to move the image.\n"
|
||
"Use the mouse wheel to zoom the image.\n"
|
||
"press 'S' or 'ENTER' to save ROIs.\n"
|
||
"press 'F' to save a full standardized screenshot.\n"
|
||
"press 'R' to output only the ROI ranges, not save.\n"
|
||
"press 'c' or 'C' (with connected) to output the ROI ranges and colors, not save.\n"
|
||
"press 'Z' or 'DELETE' or 'BACKSPACE' to remove the latest ROI.\n"
|
||
"press '0' ~ '9' to resize the window.\n"
|
||
"press 'Q' or 'ESC' to quit.\n"
|
||
"The cropped images will be saved in ./dst.\n"
|
||
)
|
||
# ---------------------------------------------------------
|
||
|
||
# 初始化 dst 文件夹
|
||
dst_path = Path.cwd() / "dst"
|
||
if not dst_path.is_dir():
|
||
dst_path.mkdir()
|
||
|
||
# 搜索并连接设备
|
||
controller = get_adb_devices_list()
|
||
if not controller:
|
||
print("ADB Devices Not Found. Reading from ./src folder...")
|
||
if controller:
|
||
set_screenshot_target_side(controller)
|
||
if controller.post_connection().failed:
|
||
print(f"Failed to connect device({device_serial}).")
|
||
exit(1)
|
||
|
||
# 初始化 cv2 窗口
|
||
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
|
||
cv2.setWindowProperty(win_name, cv2.WND_PROP_TOPMOST, 1)
|
||
cv2.setMouseCallback(win_name, mouse)
|
||
|
||
cv2.namedWindow(trackbars_name, cv2.WINDOW_NORMAL)
|
||
cv2.setWindowProperty(trackbars_name, cv2.WND_PROP_TOPMOST, 1)
|
||
cv2.createTrackbar("Scale", trackbars_name, 100, 200, trackbar_change)
|
||
|
||
cropping = False
|
||
destroyedMainColors = True
|
||
while True:
|
||
if not cropping:
|
||
std_roimage = get_std_roimage()
|
||
if std_roimage is None:
|
||
break
|
||
win_roimage = Roimage(0, 0, 0, 0, std_roimage)
|
||
crop_list.clear()
|
||
draw(crop_list)
|
||
|
||
key = cv2.waitKey(0) & 0xFF
|
||
cropping = True
|
||
# q Q esc
|
||
if key in [ord("q"), ord("Q"), 27]:
|
||
cv2.destroyAllWindows()
|
||
exit()
|
||
# 0
|
||
if key == 48:
|
||
cv2.setTrackbarPos("Scale", trackbars_name, 100)
|
||
continue
|
||
# 1 ~ 5
|
||
if key in range(49, 54):
|
||
cv2.setTrackbarPos("Scale", trackbars_name, 100 + 15 * (key - 48))
|
||
continue
|
||
# 6 ~ 9
|
||
if key in range(54, 58):
|
||
cv2.setTrackbarPos("Scale", trackbars_name, 100 - 15 * (58 - key))
|
||
continue
|
||
# z Z delete backspace
|
||
if key in [ord("z"), ord("Z"), 0, 8]:
|
||
if len(crop_list):
|
||
crop_list.pop()
|
||
continue
|
||
|
||
cropping = False
|
||
needSave = True
|
||
needColorMatch = False
|
||
connected = False
|
||
mains = []
|
||
# r R
|
||
if key in [ord("r"), ord("R")]:
|
||
needSave = False
|
||
# c C
|
||
elif key in [ord("c")]:
|
||
needSave = False
|
||
needColorMatch = True
|
||
connected = False
|
||
elif key in [ord("C")]:
|
||
needSave = False
|
||
needColorMatch = True
|
||
connected = True
|
||
# f F
|
||
elif key in [ord("f"), ord("F")]:
|
||
crop_list.append(Roimage(0, 0, 0, 0, std_roimage))
|
||
# s S enter
|
||
elif key not in [ord("s"), ord("S"), ord("\r"), ord("\n")]:
|
||
continue
|
||
|
||
for roi in crop_list:
|
||
print("")
|
||
img = roi.image
|
||
|
||
if needSave:
|
||
x1, y1, w1, h1 = roi.rectangle
|
||
x2, y2, w2, h2 = get_amplified_roi_rectangle(roi)
|
||
dst_filename: str = (
|
||
f"{file_name}_{x1}_{y1}_{w1}_{h1}__{x2}_{y2}_{w2}_{h2}.png"
|
||
)
|
||
dst_file_path = dst_path / dst_filename
|
||
print(rf"dst: {dst_file_path.absolute()}")
|
||
cv2.imwrite(str(dst_file_path), roi.image)
|
||
|
||
print(
|
||
f"original roi: {roi.rectangle}\n"
|
||
f"amplified roi: {get_amplified_roi_rectangle(roi)}"
|
||
)
|
||
|
||
if needColorMatch:
|
||
method, reverse, colors = match_color(img)
|
||
ret = {
|
||
"recognition": "ColorMatch",
|
||
"roi": [],
|
||
"method": method,
|
||
"lower": [],
|
||
"upper": [],
|
||
"count": [],
|
||
"connected": connected,
|
||
}
|
||
mainColors = []
|
||
for center, lower, upper in colors[0]:
|
||
count = colormatcher.getCount(img, lower, upper, connected, method)
|
||
ret["lower"].append(lower)
|
||
ret["upper"].append(upper)
|
||
ret["count"].append(count)
|
||
mainColor = np.zeros((80, 200, 3), dtype=np.uint8)
|
||
mainColor[:28, :] = lower
|
||
mainColor[28:58, :] = upper
|
||
mainColor[58:, :] = center
|
||
if reverse is not None:
|
||
mainColor = cv2.cvtColor(mainColor, reverse)
|
||
cv2.putText(
|
||
mainColor,
|
||
f"{lower}",
|
||
(0, 20),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.75,
|
||
(0, 0, 0),
|
||
2,
|
||
)
|
||
cv2.putText(
|
||
mainColor,
|
||
f"{upper}",
|
||
(0, 50),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.75,
|
||
(255, 255, 255),
|
||
2,
|
||
)
|
||
cv2.putText(
|
||
mainColor,
|
||
f"{count}",
|
||
(66, 76),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.75,
|
||
(127, 127, 127),
|
||
2,
|
||
)
|
||
mainColors.append(mainColor)
|
||
mains.append(cv2.hconcat(mainColors))
|
||
cv2.namedWindow("MainColors", cv2.WINDOW_NORMAL)
|
||
cv2.imshow("MainColors", cv2.vconcat(mains))
|
||
destroyedMainColors = False
|
||
print(
|
||
f"ColorMatch: {ret}".replace("'", '"')
|
||
.replace("False", "false")
|
||
.replace("True", "true")
|
||
)
|
||
elif is_window_visible("MainColors", destroyedMainColors):
|
||
cv2.destroyWindow("MainColors")
|
||
destroyedMainColors = True
|
||
|
||
print("")
|
||
|
||
print("Press 'Q'/'ESC' to quit or just continute.")
|
||
|
||
print("Exiting...")
|
||
cv2.destroyAllWindows()
|