改写战斗中替换group的逻辑

mower/utils/character_recognize.py

@@ -0,0 +1,157 @@
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+
+import cv2
+
+from mower.models import avatar, portrait
+from mower.utils import config
+from mower.utils import typealias as tp
+from mower.utils.image import cropimg, thres2
+from mower.utils.log import logger
+from mower.utils.matcher import GOOD_DISTANCE_LIMIT, flann, keypoints
+from mower.utils.vector import va
+
+
+def segment_room_select(img: tp.Image) -> list[tp.Scope]:
+    "基建房间内选干员"
+    line1 = cropimg(img, ((600, 519), (1920, 520)))
+    hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
+    mask = cv2.inRange(hsv, (98, 140, 200), (102, 255, 255))
+    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
+    line1[mask > 0] = (255,)
+    line1 = thres2(line1, 140)
+
+    last_line = line1[-1]
+    prev = last_line[0]
+    start = None
+    name_x = []
+    for i in range(1, line1.shape[1]):
+        curr = last_line[i]
+        if prev == 0 and curr == 255 and start and i - start > 186:
+            name_x.append((i + 415, i + 585))
+        elif prev == 255 and curr == 0:
+            start = i
+        prev = curr
+
+    avatar_y = ((205, 320), (625, 740))
+    avatar_p = []
+    for x in name_x:
+        for y in avatar_y:
+            avatar_p.append(tuple(zip(x, y)))
+
+    logger.debug(avatar_p)
+    return avatar_p
+
+
+def segment_team(img: tp.Image) -> list[tp.Scope]:
+    "编队界面"
+    # TODO: 利用img判断编队缺人的情况
+    result = []
+    for i in range(6):
+        x = 283 + 232 * i
+        for y in [204, 637]:
+            result.append(((x, y), va((x, y), (180, 110))))
+    logger.debug(result)
+    return result
+
+
+def segment_team_select(img: tp.Image) -> list[tp.Scope]:
+    "作战编队和训练位选干员"
+    line1 = cropimg(img, ((600, 510), (1920, 511)))
+    hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
+    mask = cv2.inRange(hsv, (98, 140, 200), (102, 255, 255))
+    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
+    line1[mask > 0] = (255,)
+    line1 = thres2(line1, 140)
+    last_line = line1[-1]
+    prev = last_line[0]
+    start = None
+    name_x = []
+    for i in range(1, line1.shape[1]):
+        curr = last_line[i]
+        if prev == 0 and curr == 255 and start and i - start > 96:
+            name_x.append((i + 415, i + 590))
+        elif prev == 255 and curr == 0:
+            start = i
+        prev = curr
+
+    avatar_y = ((200, 310), (620, 725))
+    avatar_p = []
+    for x in name_x:
+        for y in avatar_y:
+            avatar_p.append(tuple(zip(x, y)))
+
+    logger.debug(avatar_p)
+    return avatar_p
+
+
+p_list = []
+for i, img_list in portrait.items():
+    p_list.extend(img_list)
+
+a_list = []
+for i, img_list in avatar.items():
+    a_list.extend(img_list)
+
+
+def match_operator(
+    gray: tp.GrayImage, segment: list[tp.Scope], model: dict, d_list: list
+) -> tuple[tuple[str, tp.Scope]]:
+    avatar_p = {}
+    for av in segment:
+        avatar_p[av] = None, 0
+
+    kp, des = keypoints(gray)
+
+    match = partial(flann.knnMatch, trainDescriptors=des, k=2)
+
+    with ThreadPoolExecutor(max_workers=config.conf.max_workers) as executor:
+        result = executor.map(match, d_list)
+
+    for i, img_list in model.items():
+        for d1 in img_list:
+            matches = next(result)
+            good = {}
+            for j in avatar_p:
+                good[j] = []
+            for pair in matches:
+                if len(pair) != 2:
+                    continue
+                x, y = pair
+                if x.distance >= GOOD_DISTANCE_LIMIT * y.distance:
+                    continue
+                for j in avatar_p:
+                    kpj = kp[x.trainIdx]
+                    if j[0][0] < kpj.pt[0] < j[1][0] and j[0][1] < kpj.pt[1] < j[1][1]:
+                        good[j].append(x)
+            for j, g in good.items():
+                score = len(g) / len(d1)
+                if avatar_p[j][1] < score:
+                    avatar_p[j] = i, score
+
+    op_name = [p[0] for p in avatar_p.values()]
+    logger.debug(op_name)
+
+    return tuple(zip(op_name, segment))
+
+
+match_portrait = partial(match_operator, model=portrait, d_list=p_list)
+match_avatar = partial(match_operator, model=avatar, d_list=a_list)
+
+
+def operator_room_select(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
+    "基建房间内选干员"
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    return match_portrait(gray, segment_room_select(img))
+
+
+def operator_team(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
+    "编队界面"
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    return match_portrait(gray, segment_team(img))
+
+
+def operator_team_select(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
+    "作战编队和训练位选干员"
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    return match_portrait(gray, segment_team_select(img))