Konano/arknights-mower

View on GitHub
arknights_mower/utils/recognize.py

Summary

Maintainability
B
6 hrs
Test Coverage
from __future__ import annotations

import time
from typing import List, Optional

import cv2
import numpy as np

from .. import __rootdir__
from . import config, detector
from . import typealias as tp
from .device import Device
from .image import bytes2img, cropimg, loadimg, thres2
from .log import logger, save_screenshot
from .matcher import Matcher
from .scene import Scene, SceneComment


class RecognizeError(Exception):
    pass


class Recognizer(object):

    def __init__(self, device: Device, screencap: bytes = None) -> None:
        self.device = device
        self.start(screencap)

    def start(self, screencap: bytes = None, build: bool = True) -> None:
        """ init with screencap, build matcher  """
        retry_times = config.MAX_RETRYTIME
        while retry_times > 0:
            try:
                if screencap is not None:
                    self.screencap = screencap
                else:
                    self.screencap = self.device.screencap()
                self.img = bytes2img(self.screencap, False)
                self.gray = bytes2img(self.screencap, True)
                self.h, self.w, _ = self.img.shape
                self.matcher = Matcher(self.gray) if build else None
                self.scene = Scene.UNDEFINED
                return
            except cv2.error as e:
                logger.warning(e)
                retry_times -= 1
                time.sleep(1)
                continue
        raise RuntimeError('init Recognizer failed')

    def update(self, screencap: bytes = None, rebuild: bool = True) -> None:
        """ rebuild matcher """
        self.start(screencap, rebuild)

    def color(self, x: int, y: int) -> tp.Pixel:
        """ get the color of the pixel """
        return self.img[y][x]

    def save_screencap(self, folder):
        save_screenshot(self.screencap, subdir=f'{folder}/{self.h}x{self.w}')

    def get_scene(self) -> int:
        """ get the current scene in the game """
        if self.scene != Scene.UNDEFINED:
            return self.scene
        if self.find('connecting', scope=((self.w//2, self.h//10*8), (self.w//4*3, self.h))) is not None:
            self.scene = Scene.CONNECTING
        elif self.find('index_nav', thres=250, scope=((0, 0), (100+self.w//4, self.h//10))) is not None:
            self.scene = Scene.INDEX
        elif self.find('nav_index') is not None:
            self.scene = Scene.NAVIGATION_BAR
        elif self.find('close_mine') is not None:
            self.scene = Scene.CLOSE_MINE
        elif self.find('materiel_ico') is not None:
            self.scene = Scene.MATERIEL
        elif self.find('read_mail') is not None:
            self.scene = Scene.MAIL
        elif self.find('loading') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading2') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading3') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading4') is not None:
            self.scene = Scene.LOADING
        elif self.is_black():
            self.scene = Scene.LOADING
        elif self.find('ope_plan') is not None:
            self.scene = Scene.OPERATOR_BEFORE
        elif self.find('ope_select_start') is not None:
            self.scene = Scene.OPERATOR_SELECT
        elif self.find('ope_agency_going') is not None:
            self.scene = Scene.OPERATOR_ONGOING
        elif self.find('ope_elimi_finished') is not None:
            self.scene = Scene.OPERATOR_ELIMINATE_FINISH
        elif self.find('ope_finish') is not None:
            self.scene = Scene.OPERATOR_FINISH
        elif self.find('ope_recover_potion_on') is not None:
            self.scene = Scene.OPERATOR_RECOVER_POTION
        elif self.find('ope_recover_originite_on') is not None:
            self.scene = Scene.OPERATOR_RECOVER_ORIGINITE
        elif self.find('double_confirm') is not None:
            if self.find('network_check') is not None:
                self.scene = Scene.NETWORK_CHECK
            else:
                self.scene = Scene.DOUBLE_CONFIRM
        elif self.find('ope_firstdrop') is not None:
            self.scene = Scene.OPERATOR_DROP
        elif self.find('ope_eliminate') is not None:
            self.scene = Scene.OPERATOR_ELIMINATE
        elif self.find('ope_elimi_agency_panel') is not None:
            self.scene = Scene.OPERATOR_ELIMINATE_AGENCY
        elif self.find('ope_giveup') is not None:
            self.scene = Scene.OPERATOR_GIVEUP
        elif self.find('ope_failed') is not None:
            self.scene = Scene.OPERATOR_FAILED
        elif self.find('friend_list_on') is not None:
            self.scene = Scene.FRIEND_LIST_ON
        elif self.find('credit_visiting') is not None:
            self.scene = Scene.FRIEND_VISITING
        elif self.find('infra_overview') is not None:
            self.scene = Scene.INFRA_MAIN
        elif self.find('infra_todo') is not None:
            self.scene = Scene.INFRA_TODOLIST
        elif self.find('clue') is not None:
            self.scene = Scene.INFRA_CONFIDENTIAL
        elif self.find('arrange_check_in') or self.find('arrange_check_in_on') is not None:
            self.scene = Scene.INFRA_DETAILS
        elif self.find('infra_overview_in') is not None:
            self.scene = Scene.INFRA_ARRANGE
        elif self.find('arrange_confirm') is not None:
            self.scene = Scene.INFRA_ARRANGE_CONFIRM
        elif self.find('friend_list') is not None:
            self.scene = Scene.FRIEND_LIST_OFF
        elif self.find("mission_trainee_on") is not None:
            self.scene = Scene.MISSION_TRAINEE
        elif self.find('mission_daily_on') is not None:
            self.scene = Scene.MISSION_DAILY
        elif self.find('mission_weekly_on') is not None:
            self.scene = Scene.MISSION_WEEKLY
        elif self.find('terminal_pre') is not None:
            self.scene = Scene.TERMINAL_MAIN
        elif self.find('open_recruitment') is not None:
            self.scene = Scene.RECRUIT_MAIN
        elif self.find('recruiting_instructions') is not None:
            self.scene = Scene.RECRUIT_TAGS
        elif self.find('agent_token') is not None:
            self.scene = Scene.RECRUIT_AGENT
        elif self.find('agent_token_1080_1440') is not None:
            self.scene = Scene.RECRUIT_AGENT
        elif self.find('agent_token_900_1440') is not None:
            self.scene = Scene.RECRUIT_AGENT
        elif self.find('agent_unlock') is not None:
            self.scene = Scene.SHOP_CREDIT
        elif self.find('shop_credit_2') is not None:
            self.scene = Scene.SHOP_OTHERS
        elif self.find('shop_cart') is not None:
            self.scene = Scene.SHOP_CREDIT_CONFIRM
        elif self.find('shop_assist') is not None:
            self.scene = Scene.SHOP_ASSIST
        elif self.find('login_logo') is not None and self.find('hypergryph') is not None:
            if self.find('login_awake') is not None:
                self.scene = Scene.LOGIN_QUICKLY
            elif self.find('login_account') is not None:
                self.scene = Scene.LOGIN_MAIN
            elif self.find('login_iknow') is not None:
                self.scene = Scene.LOGIN_ANNOUNCE
            else:
                self.scene = Scene.LOGIN_MAIN_NOENTRY
        elif self.find('register') is not None:
            self.scene = Scene.LOGIN_REGISTER
        elif self.find('login_loading') is not None:
            self.scene = Scene.LOGIN_LOADING
        elif self.find('login_iknow') is not None:
            self.scene = Scene.LOGIN_ANNOUNCE
        elif self.find('12cadpa') is not None:
            if self.find('cadpa_detail') is not None:
                self.scene = Scene.LOGIN_CADPA_DETAIL
            else:
                self.scene = Scene.LOGIN_START
        elif detector.announcement_close(self.img) is not None:
            self.scene = Scene.ANNOUNCEMENT
        elif self.find('skip') is not None:
            self.scene = Scene.SKIP
        elif self.find('upgrade') is not None:
            self.scene = Scene.UPGRADE
        elif detector.confirm(self.img) is not None:
            self.scene = Scene.CONFIRM
        elif self.find('login_verify') is not None:
            self.scene = Scene.LOGIN_INPUT
        elif self.find('login_captcha') is not None:
            self.scene = Scene.LOGIN_CAPTCHA
        elif self.find('login_connecting') is not None:
            self.scene = Scene.LOGIN_LOADING
        elif self.find('main_theme') is not None:
            self.scene = Scene.TERMINAL_MAIN_THEME
        elif self.find('episode') is not None:
            self.scene = Scene.TERMINAL_EPISODE
        elif self.find('biography') is not None:
            self.scene = Scene.TERMINAL_BIOGRAPHY
        elif self.find('collection') is not None:
            self.scene = Scene.TERMINAL_COLLECTION
        elif self.find('login_bilibili') is not None:
            self.scene = Scene.LOGIN_BILIBILI
        elif self.find('loading6') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading7') is not None:
            self.scene = Scene.LOADING
        elif self.find('arrange_order_options_scene') is not None:
            self.scene = Scene.INFRA_ARRANGE_ORDER
        else:
            self.scene = Scene.UNKNOWN
            self.device.check_current_focus()
        # save screencap to analyse
        if config.SCREENSHOT_PATH is not None:
            self.save_screencap(self.scene)
        logger.info(f'Scene: {self.scene}: {SceneComment[self.scene]}')
        return self.scene

    def get_infra_scene(self)-> int:
        if self.scene != Scene.UNDEFINED:
            return self.scene
        if self.find('connecting', scope=((self.w//2, self.h//10*8), (self.w//4*3, self.h))) is not None:
            self.scene = Scene.CONNECTING
        elif self.find('double_confirm') is not None:
            if self.find('network_check') is not None:
                self.scene = Scene.NETWORK_CHECK
            else:
                self.scene = Scene.DOUBLE_CONFIRM
        elif self.find('infra_overview') is not None:
            self.scene = Scene.INFRA_MAIN
        elif self.find('infra_todo') is not None:
            self.scene = Scene.INFRA_TODOLIST
        elif self.find('clue') is not None:
            self.scene = Scene.INFRA_CONFIDENTIAL
        elif self.find('arrange_check_in') or self.find('arrange_check_in_on') is not None:
            self.scene = Scene.INFRA_DETAILS
        elif self.find('infra_overview_in') is not None:
            self.scene = Scene.INFRA_ARRANGE
        elif self.find('arrange_confirm') is not None:
            self.scene = Scene.INFRA_ARRANGE_CONFIRM
        elif self.find('arrange_order_options_scene') is not None:
            self.scene = Scene.INFRA_ARRANGE_ORDER
        elif self.find('loading') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading2') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading3') is not None:
            self.scene = Scene.LOADING
        elif self.find('loading4') is not None:
            self.scene = Scene.LOADING
        elif self.find('index_nav', thres=250, scope=((0, 0), (100+self.w//4, self.h//10))) is not None:
            self.scene = Scene.INDEX
        elif self.is_black():
            self.scene = Scene.LOADING
        else:
            self.scene = Scene.UNKNOWN
            self.device.check_current_focus()
        # save screencap to analyse
        if config.SCREENSHOT_PATH is not None:
            self.save_screencap(self.scene)
        logger.info(f'Scene: {self.scene}: {SceneComment[self.scene]}')
        return self.scene

    def is_black(self) -> None:
        """ check if the current scene is all black """
        return np.max(self.gray[:, 105:-105]) < 16

    def nav_button(self):
        """ find navigation button """
        return self.find('nav_button', thres=128, scope=((0, 0), (100+self.w//4, self.h//10)))

    def find(self, res: str, draw: bool = False, scope: tp.Scope = None, thres: int = None, judge: bool = True, strict: bool = False) -> tp.Scope:
        """
        查找元素是否出现在画面中

        :param res: 待识别元素资源文件名
        :param draw: 是否将识别结果输出到屏幕
        :param scope: ((x0, y0), (x1, y1)),提前限定元素可能出现的范围
        :param thres: 是否在匹配前对图像进行二值化处理
        :param judge: 是否加入更加精确的判断
        :param strict: 是否启用严格模式,未找到时报错

        :return ret: 若匹配成功,则返回元素在游戏界面中出现的位置,否则返回 None
        """
        logger.debug(f'find: {res}')
        res = f'{__rootdir__}/resources/{res}.png'

        if thres is not None:
            # 对图像二值化处理
            res_img = thres2(loadimg(res, True), thres)
            gray_img = cropimg(self.gray, scope)
            matcher = Matcher(thres2(gray_img, thres))
            ret = matcher.match(res_img, draw=draw, judge=judge)
        else:
            res_img = loadimg(res, True)
            matcher = self.matcher
            ret = matcher.match(res_img, draw=draw, scope=scope, judge=judge)
        if strict and ret is None:
            raise RecognizeError(f"Can't find '{res}'") 
        return ret

    def score(self, res: str, draw: bool = False, scope: tp.Scope = None, thres: int = None) -> Optional[List[float]]:
        """
        查找元素是否出现在画面中,并返回分数

        :param res: 待识别元素资源文件名
        :param draw: 是否将识别结果输出到屏幕
        :param scope: ((x0, y0), (x1, y1)),提前限定元素可能出现的范围
        :param thres: 是否在匹配前对图像进行二值化处理

        :return ret: 若匹配成功,则返回元素在游戏界面中出现的位置,否则返回 None
        """
        logger.debug(f'find: {res}')
        res = f'{__rootdir__}/resources/{res}.png'

        if thres is not None:
            # 对图像二值化处理
            res_img = thres2(loadimg(res, True), thres)
            gray_img = cropimg(self.gray, scope)
            matcher = Matcher(thres2(gray_img, thres))
            score = matcher.score(res_img, draw=draw, only_score=True)
        else:
            res_img = loadimg(res, True)
            matcher = self.matcher
            score = matcher.score(res_img, draw=draw, scope=scope, only_score=True)
        return score