diff --git a/assets/share/combat/prepare/OCR_TRAILBLAZE_POWER.png b/assets/share/combat/prepare/OCR_TRAILBLAZE_POWER.png new file mode 100644 index 000000000..5ecfd9c55 Binary files /dev/null and b/assets/share/combat/prepare/OCR_TRAILBLAZE_POWER.png differ diff --git a/module/ocr/ocr.py b/module/ocr/ocr.py index 4c4e353ce..bd80a4a69 100644 --- a/module/ocr/ocr.py +++ b/module/ocr/ocr.py @@ -1,6 +1,7 @@ import time import cv2 +import re from ppocronnx.predict_system import BoxedResult import module.config.server as server @@ -204,3 +205,48 @@ class Ocr: logger.attr(name=f'{self.name} matched', text=results) return results + + +class Digit(Ocr): + def __init__(self, button: ButtonWrapper, lang='ch', name=None): + super().__init__(button, lang=lang, name=name) + + def after_process(self, result) -> int: + """ + Returns: + int: + """ + result = super().after_process(result) + logger.attr(name=self.name, text=str(result)) + + res = re.search(r'(\d+)', result) + if res: + return int(res.group(1)) + else: + logger.warning(f'No digit found in {result}') + return 0 + + +class DigitCounter(Ocr): + def __init__(self, button: ButtonWrapper, lang='ch', name=None): + super().__init__(button, lang=lang, name=name) + + def after_process(self, result) -> tuple[int, int, int]: + """ + Do OCR on a counter, such as `14/15`, and returns 14, 1, 15 + + Returns: + int: + """ + result = super().after_process(result) + logger.attr(name=self.name, text=str(result)) + + res = re.search(r'(\d+)/(\d+)', result) + if res: + groups = [int(s) for s in res.groups()] + current, total = int(groups[0]), int(groups[1]) + current = min(current, total) + return current, total - current, total + else: + logger.warning(f'No digit counter found in {result}') + return 0, 0, 0 diff --git a/tasks/base/state.py b/tasks/base/state.py new file mode 100644 index 000000000..fe44ac8f0 --- /dev/null +++ b/tasks/base/state.py @@ -0,0 +1,45 @@ +import datetime + +from pydantic import BaseModel + +from module.config.utils import DEFAULT_TIME +from module.logger import logger + + +def now(): + return datetime.datetime.now().replace(microsecond=0) + + +class StateValue(BaseModel): + value: int = 0 + time: datetime.datetime = DEFAULT_TIME + + +TrailblazePowerMax = 180 +ImmersifierMax = 8 + + +class StateStorage(BaseModel): + TrailblazePower = StateValue() + Immersifier = StateValue() + + def __setattr__(self, key, value): + if key in super().__getattribute__('__fields__'): + storage = super().__getattribute__(key) + storage.value = value + storage.time = now() + else: + super().__setattr__(key, value) + + def __getattribute__(self, item): + if item in super().__getattribute__('__fields__'): + storage = super().__getattribute__(item) + if storage.time == DEFAULT_TIME: + logger.warning(f'Trying to get state {item} but it is never set') + return storage.value + else: + return super().__getattribute__(item) + + +class StateMixin: + state = StateStorage() diff --git a/tasks/base/ui.py b/tasks/base/ui.py index 15f3c61b2..343697533 100644 --- a/tasks/base/ui.py +++ b/tasks/base/ui.py @@ -6,10 +6,11 @@ from module.ocr.ocr import Ocr from tasks.base.assets.assets_base_page import CLOSE from tasks.base.page import Page, page_main from tasks.base.popup import PopupHandler +from tasks.base.state import StateMixin from tasks.combat.assets.assets_combat_prepare import COMBAT_PREPARE -class UI(PopupHandler): +class UI(PopupHandler, StateMixin): ui_current: Page def ui_page_appear(self, page): diff --git a/tasks/combat/assets/assets_combat_prepare.py b/tasks/combat/assets/assets_combat_prepare.py index d5d8e3942..86100a046 100644 --- a/tasks/combat/assets/assets_combat_prepare.py +++ b/tasks/combat/assets/assets_combat_prepare.py @@ -13,6 +13,16 @@ COMBAT_PREPARE = ButtonWrapper( button=(956, 640, 1224, 676), ), ) +OCR_TRAILBLAZE_POWER = ButtonWrapper( + name='OCR_TRAILBLAZE_POWER', + share=Button( + file='./assets/share/combat/prepare/OCR_TRAILBLAZE_POWER.png', + area=(1043, 26, 1131, 48), + search=(1023, 6, 1151, 68), + color=(43, 46, 53), + button=(1043, 26, 1131, 48), + ), +) OCR_WAVE_COUNT = ButtonWrapper( name='OCR_WAVE_COUNT', share=Button( diff --git a/tasks/combat/prepare.py b/tasks/combat/prepare.py index e9e641807..ddbefbc0e 100644 --- a/tasks/combat/prepare.py +++ b/tasks/combat/prepare.py @@ -1,31 +1,11 @@ -import re - -from module.logger import logger -from module.ocr.ocr import Ocr +from module.ocr.ocr import Digit, DigitCounter from tasks.base.ui import UI -from tasks.combat.assets.assets_combat_prepare import OCR_WAVE_COUNT, WAVE_MINUS, WAVE_PLUS - - -class WaveCount(Ocr): - def after_process(self, result): - """ - Returns: - int: - """ - result = super().after_process(result) - logger.attr(name=self.name, text=str(result)) - - res = re.search(r'(\d)', result) - if res: - result = int(res.group(1)) - if 1 <= result <= 6: - return result - else: - logger.warning(f'Unexpected combat wave: {result}') - return 0 - else: - logger.warning('Cannot find wave count') - return 0 +from tasks.combat.assets.assets_combat_prepare import ( + OCR_TRAILBLAZE_POWER, + OCR_WAVE_COUNT, + WAVE_MINUS, + WAVE_PLUS +) class CombatPrepare(UI): @@ -33,9 +13,21 @@ class CombatPrepare(UI): """ Args: count: 1 to 6 + + Pages: + in: COMBAT_PREPARE """ self.ui_ensure_index( - count, letter=WaveCount(OCR_WAVE_COUNT), + count, letter=Digit(OCR_WAVE_COUNT), next_button=WAVE_PLUS, prev_button=WAVE_MINUS, skip_first_screenshot=True ) + + def combat_get_trailblaze_power(self) -> int: + """ + Pages: + in: COMBAT_PREPARE or COMBAT_REPEAT + """ + current, _, _ = DigitCounter(OCR_TRAILBLAZE_POWER).ocr_single_line(self.device.image) + self.state.TrailblazePower = current + return current