159 lines
4.8 KiB
Python
159 lines
4.8 KiB
Python
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 in_scope, 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, ((540, 513), (1920, 514)))
|
|
hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
|
|
mask = cv2.inRange(hsv, (98, 150, 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 == 255 and curr == 0 and start and i - start > 90:
|
|
name_x.append((i + 450, i + 630))
|
|
elif prev == 0 and curr == 255:
|
|
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
|
|
|
|
|
|
def match_operator(
|
|
gray: tp.GrayImage, segment: list[tp.Scope], model: list
|
|
) -> tuple[tuple[str, tp.Scope]]:
|
|
if len(segment) == 0:
|
|
logger.debug(f"{len(segment)=}")
|
|
return ()
|
|
|
|
avatar_p = [[None, 0]] * len(segment)
|
|
|
|
left = min(s[0][0] for s in segment)
|
|
right = max(s[1][0] for s in segment)
|
|
top = min(s[0][1] for s in segment)
|
|
bottom = max(s[1][1] for s in segment)
|
|
gray = cropimg(gray, ((left, top), (right, bottom)))
|
|
ori_kp, ori_des = keypoints(gray)
|
|
|
|
def match(item):
|
|
name, des = item
|
|
return name, flann.knnMatch(des, ori_des, k=2)
|
|
|
|
if config.conf.max_workers > 1:
|
|
with ThreadPoolExecutor(max_workers=config.conf.max_workers) as executor:
|
|
result = executor.map(match, model)
|
|
else:
|
|
result = [match(i) for i in model]
|
|
|
|
for name, matches in result:
|
|
good = []
|
|
for pair in matches:
|
|
if len(pair) == 2:
|
|
x, y = pair
|
|
if x.distance < GOOD_DISTANCE_LIMIT * y.distance:
|
|
good.append(ori_kp[x.trainIdx].pt)
|
|
elif len(pair) == 1:
|
|
x = pair[0]
|
|
good.append(ori_kp[x.trainIdx].pt)
|
|
good = [va(g, (left, top)) for g in good]
|
|
count = [0] * len(segment)
|
|
for i, scope in enumerate(segment):
|
|
for g in good:
|
|
if in_scope(scope, g):
|
|
count[i] += 1
|
|
for i, c in enumerate(count):
|
|
if c > avatar_p[i][1]:
|
|
avatar_p[i] = name, c
|
|
|
|
op_name = [p[0] for p in avatar_p]
|
|
logger.debug(op_name)
|
|
|
|
return tuple(zip(op_name, segment))
|
|
|
|
|
|
match_portrait = partial(match_operator, model=portrait)
|
|
match_avatar = partial(match_operator, model=avatar)
|
|
|
|
|
|
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))
|