Add: Get trailblaze power before and after combat

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

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()

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):

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(

tasks/combat/prepare.py

@@ -1,31 +1,11 @@
-import re
+from module.ocr.ocr import Digit, DigitCounter
-
-from module.logger import logger
-from module.ocr.ocr import Ocr
 from tasks.base.ui import UI
-from tasks.combat.assets.assets_combat_prepare import OCR_WAVE_COUNT, WAVE_MINUS, WAVE_PLUS
+from tasks.combat.assets.assets_combat_prepare import (
-
+    OCR_TRAILBLAZE_POWER,
-
+    OCR_WAVE_COUNT,
-class WaveCount(Ocr):
+    WAVE_MINUS,
-    def after_process(self, result):
+    WAVE_PLUS
-        """
+)
-        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
 
 
 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