212 lines
6.7 KiB
Python
212 lines
6.7 KiB
Python
from functools import partial
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from mower.models import avatar, portrait
|
|
from mower.utils import typealias as tp
|
|
from mower.utils.image import cmatch, cropimg, loadres, thres2
|
|
from mower.utils.log import logger
|
|
from mower.utils.matcher import GOOD_DISTANCE_LIMIT, flann, search_params
|
|
from mower.utils.vector import in_scope, va
|
|
|
|
FAST = cv2.FastFeatureDetector_create()
|
|
ORB = cv2.ORB_create(nfeatures=1000000, nlevels=1)
|
|
|
|
FLANN_INDEX_LSH = 6
|
|
index_params = dict(
|
|
algorithm=FLANN_INDEX_LSH,
|
|
table_number=1,
|
|
key_size=18,
|
|
multi_probe_level=1,
|
|
)
|
|
flann1 = cv2.FlannBasedMatcher(index_params, search_params)
|
|
|
|
|
|
def fast_keypoints(img: tp.GrayImage, mask: tp.GrayImage | None = None):
|
|
return ORB.compute(img, FAST.detect(img, mask))
|
|
|
|
|
|
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, (95, 230, 200), (105, 255, 230))
|
|
line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
|
|
line1[mask > 0] = (255,)
|
|
mask = cv2.inRange(hsv, (10, 0, 0), (20, 255, 255))
|
|
line1[mask > 0] = (0,)
|
|
line1 = thres2(line1, 100)
|
|
line = line1[-1]
|
|
prev = line[0]
|
|
start = None
|
|
name_x = []
|
|
for i in range(1, line1.shape[1]):
|
|
curr = line[i]
|
|
if prev == 0 and curr == 255 and start and i - start > 186:
|
|
name_x.append((i + 434, i + 574))
|
|
elif prev == 255 and curr == 0:
|
|
start = i
|
|
prev = curr
|
|
|
|
avatar_y = ((242, 317), (662, 737))
|
|
avatar_p = []
|
|
for x in name_x:
|
|
for y in avatar_y:
|
|
avatar_p.append(tuple(zip(x, y)))
|
|
|
|
empty = loadres("choose_agent/empty")
|
|
if avatar_p and cmatch(cropimg(img, avatar_p[-1]), empty, 20):
|
|
avatar_p.pop()
|
|
|
|
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, ((545, 510), (1920, 511)))
|
|
# 选中时的蓝色边框变白，信赖排序时的橙色变黑
|
|
hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
|
|
mask = cv2.inRange(hsv, (95, 230, 200), (105, 255, 230))
|
|
line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
|
|
line1[mask > 0] = (255,)
|
|
mask = cv2.inRange(hsv, (10, 0, 0), (20, 255, 255))
|
|
line1[mask > 0] = (0,)
|
|
line1 = thres2(line1, 100)
|
|
line = line1[-1]
|
|
prev = line[0]
|
|
start = None
|
|
name_x = []
|
|
for i in range(1, line1.shape[1]):
|
|
curr = line[i]
|
|
if prev == 0 and curr == 255 and start and i - start > 90:
|
|
name_x.append((i + 381, i + 521))
|
|
elif prev == 255 and curr == 0:
|
|
start = i
|
|
prev = curr
|
|
|
|
avatar_y = ((229, 304), (645, 720))
|
|
avatar_p = []
|
|
for x in name_x:
|
|
for y in avatar_y:
|
|
avatar_p.append(tuple(zip(x, y)))
|
|
|
|
empty = loadres("choose_agent/empty")
|
|
if avatar_p and cmatch(cropimg(img, avatar_p[-1]), empty, 20):
|
|
avatar_p.pop()
|
|
|
|
logger.debug(avatar_p)
|
|
return avatar_p
|
|
|
|
|
|
def match_operator(
|
|
gray: tp.GrayImage, segment: list[tp.Scope], model: list
|
|
) -> list[tuple[str | None, tp.Scope]]:
|
|
if len(segment) == 0:
|
|
logger.debug(f"{len(segment)=}")
|
|
return ()
|
|
|
|
w = model["w"]
|
|
h = model["h"]
|
|
col = model["col"]
|
|
name_list = model["name"]
|
|
pt2 = model["pt"]
|
|
des2 = model["des"]
|
|
|
|
mask = np.zeros((1080, 1920), dtype=np.uint8)
|
|
for seg in segment:
|
|
cv2.rectangle(mask, seg[0], seg[1], (255,), -1)
|
|
kp1, des1 = fast_keypoints(gray, mask)
|
|
|
|
result = [{} for _ in range(len(segment))]
|
|
for match in flann.knnMatch(des1, des2, k=2):
|
|
if len(match) != 2:
|
|
continue
|
|
m, n = match
|
|
if m.distance > GOOD_DISTANCE_LIMIT * n.distance:
|
|
continue
|
|
x, y = pt2[m.trainIdx]
|
|
x, y = int(x / w), int(y / h)
|
|
name = name_list[y * col + x]
|
|
for i, seg in enumerate(segment):
|
|
if in_scope(seg, kp1[m.queryIdx].pt):
|
|
if name in result[i]:
|
|
result[i][name] += 1
|
|
else:
|
|
result[i][name] = 1
|
|
break
|
|
result = [list(r.items()) for r in result]
|
|
result = [max(r, key=lambda p: p[1]) if r else None for r in result]
|
|
logger.debug(result)
|
|
|
|
bad_index = []
|
|
for i, r in enumerate(result):
|
|
if r is None or r[1] < 10:
|
|
bad_index.append(i)
|
|
if bad_index:
|
|
mask = np.zeros((1080, 1920), dtype=np.uint8)
|
|
for i in bad_index:
|
|
cv2.rectangle(mask, segment[i][0], segment[i][1], (255,), -1)
|
|
kp1, des1 = fast_keypoints(gray, mask)
|
|
|
|
result1 = [{} for _ in range(len(bad_index))]
|
|
for match in flann1.knnMatch(des1, des2, k=2):
|
|
if len(match) != 2:
|
|
continue
|
|
m, n = match
|
|
if m.distance > GOOD_DISTANCE_LIMIT * n.distance:
|
|
continue
|
|
x, y = pt2[m.trainIdx]
|
|
x, y = int(x / w), int(y / h)
|
|
name = name_list[y * col + x]
|
|
for i, j in enumerate(bad_index):
|
|
if in_scope(segment[j], kp1[m.queryIdx].pt):
|
|
if name in result1[i]:
|
|
result1[i][name] += 1
|
|
else:
|
|
result1[i][name] = 1
|
|
break
|
|
result1 = [list(r.items()) for r in result1]
|
|
result1 = [max(r, key=lambda p: p[1]) if r else None for r in result1]
|
|
logger.debug(result1)
|
|
|
|
for i, r in enumerate(result1):
|
|
result[bad_index[i]] = r
|
|
|
|
return list(zip([r[0] if r else None for r in result], segment))
|
|
|
|
|
|
match_portrait = partial(match_operator, model=portrait)
|
|
match_avatar = partial(match_operator, model=avatar)
|
|
|
|
|
|
def operator_room_select(img: tp.Image) -> list[tuple[str | None, tp.Scope]]:
|
|
"基建房间内选干员"
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
return match_portrait(gray, segment_room_select(img))
|
|
|
|
|
|
def operator_team(img: tp.Image) -> list[tuple[str | None, tp.Scope]]:
|
|
"编队界面"
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
return match_portrait(gray, segment_team(img))
|
|
|
|
|
|
def operator_team_select(img: tp.Image) -> list[tuple[str | None, tp.Scope]]:
|
|
"作战编队选干员"
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
return match_portrait(gray, segment_team_select(img))
|