Add: Get item amount in COMBAT_PREPARE

tasks/dungeon/assets/assets_dungeon_obtain.py

@@ -0,0 +1,65 @@
+from module.base.button import Button, ButtonWrapper
+
+# This file was auto-generated, do not modify it manually. To generate:
+# ``` python -m dev_tools.button_extract ```
+
+ITEM_AMOUNT = ButtonWrapper(
+    name='ITEM_AMOUNT',
+    share=Button(
+        file='./assets/share/dungeon/obtain/ITEM_AMOUNT.png',
+        area=(190, 521, 490, 539),
+        search=(170, 501, 510, 559),
+        color=(195, 190, 188),
+        button=(190, 521, 490, 539),
+    ),
+)
+ITEM_CLOSE = ButtonWrapper(
+    name='ITEM_CLOSE',
+    share=Button(
+        file='./assets/share/dungeon/obtain/ITEM_CLOSE.png',
+        area=(1043, 185, 1073, 215),
+        search=(1023, 165, 1093, 235),
+        color=(170, 170, 170),
+        button=(1043, 185, 1073, 215),
+    ),
+)
+ITEM_NAME = ButtonWrapper(
+    name='ITEM_NAME',
+    share=Button(
+        file='./assets/share/dungeon/obtain/ITEM_NAME.png',
+        area=(495, 187, 855, 211),
+        search=(475, 167, 875, 231),
+        color=(176, 177, 179),
+        button=(495, 187, 855, 211),
+    ),
+)
+OBTAIN_1 = ButtonWrapper(
+    name='OBTAIN_1',
+    share=Button(
+        file='./assets/share/dungeon/obtain/OBTAIN_1.png',
+        area=(813, 414, 877, 478),
+        search=(793, 394, 897, 498),
+        color=(118, 96, 131),
+        button=(813, 414, 877, 478),
+    ),
+)
+OBTAIN_2 = ButtonWrapper(
+    name='OBTAIN_2',
+    share=Button(
+        file='./assets/share/dungeon/obtain/OBTAIN_2.png',
+        area=(889, 414, 953, 478),
+        search=(869, 394, 973, 498),
+        color=(71, 96, 145),
+        button=(889, 414, 953, 478),
+    ),
+)
+OBTAIN_3 = ButtonWrapper(
+    name='OBTAIN_3',
+    share=Button(
+        file='./assets/share/dungeon/obtain/OBTAIN_3.png',
+        area=(965, 414, 1029, 478),
+        search=(945, 394, 1049, 498),
+        color=(76, 101, 109),
+        button=(965, 414, 1029, 478),
+    ),
+)

tasks/dungeon/obtain.py

@@ -0,0 +1,222 @@
+import re
+
+from pydantic import BaseModel
+
+from module.base.timer import Timer
+from module.exception import ScriptError
+from module.logger import logger
+from module.ocr.ocr import Digit
+from tasks.base.ui import UI
+from tasks.combat.assets.assets_combat_prepare import COMBAT_PREPARE
+from tasks.dungeon.assets.assets_dungeon_obtain import *
+from tasks.dungeon.keywords import DungeonList
+from tasks.planner.keywords import ITEM_CLASSES
+from tasks.planner.keywords.classes import ItemBase
+from tasks.planner.result import OcrItemName
+
+
+class ItemAmount(BaseModel):
+    item: ItemBase
+    amount: int
+
+
+class OcrItemAmount(Digit):
+    def format_result(self, result):
+        res = re.split(r'[:：;；]', result)
+        result = res[-1]
+        return super().format_result(result)
+
+
+class DungeonObtain(UI):
+    """
+    Parse items that can be obtained from dungeon
+
+    Pages:
+        in: COMBAT_PREPARE
+    """
+
+    def _obtain_enter(self, entry, skip_first_screenshot=True):
+        """
+        Args:
+            entry: Item entry
+            skip_first_screenshot:
+
+        Pages:
+            in: COMBAT_PREPARE
+            out: ITEM_CLOSE
+        """
+        logger.info(f'Obtain enter {entry}')
+        self.interval_clear(COMBAT_PREPARE)
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if self.appear(ITEM_CLOSE):
+                break
+            if self.appear(COMBAT_PREPARE, interval=2):
+                self.device.click(entry)
+                self.interval_reset(COMBAT_PREPARE)
+                continue
+
+        # Wait animation
+        timeout = Timer(1.4, count=7).start()
+        skip_first_screenshot = True
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if self.image_color_count(ITEM_NAME, color=(0, 0, 0), threshold=221, count=50):
+                break
+            if timeout.reached():
+                logger.warning('Wait obtain item timeout')
+                break
+
+    def _obtain_close(self, skip_first_screenshot=True):
+        """
+        Args:
+            skip_first_screenshot:
+
+        Pages:
+            in: ITEM_CLOSE
+            out: COMBAT_PREPARE
+        """
+        logger.info(f'Obtain close')
+        self.interval_clear(ITEM_CLOSE)
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if self.appear(COMBAT_PREPARE):
+                break
+            if self.appear_then_click(ITEM_CLOSE, interval=2):
+                continue
+
+    @staticmethod
+    def _obtain_get_entry(dungeon: DungeonList, index: int = 1, prev: ItemAmount = None):
+        """
+        Args:
+            dungeon: Current dungeon
+            index: 1 to 3, index to check
+            prev: Previous item checked
+
+        Returns:
+            ButtonWrapper: Item entry, or None if no more check needed
+        """
+        if (index > 1 and prev is None) or (index <= 1 and prev is not None):
+            raise ScriptError(f'_obtain_get_entry: index and prev must be set together, index={index}, prev={prev}')
+
+        if index > 3:
+            return None
+
+        def may_obtain_one():
+            if prev is None:
+                return OBTAIN_1
+            else:
+                return None
+
+        def may_obtain_multi():
+            if prev is None:
+                return OBTAIN_1
+            # End at the item with the lowest rarity
+            if prev.item.is_rarity_green:
+                return None
+            if index == 2:
+                return OBTAIN_2
+            if index == 3:
+                return OBTAIN_3
+
+        if dungeon is None:
+            return may_obtain_multi()
+        if dungeon.is_Echo_of_War:
+            return may_obtain_one()
+        if dungeon.is_Cavern_of_Corrosion:
+            return None
+        if dungeon.is_Stagnant_Shadow:
+            return may_obtain_one()
+        if dungeon.is_Calyx_Golden:
+            if dungeon.is_Calyx_Golden_Treasures:
+                return may_obtain_one()
+            else:
+                return may_obtain_multi()
+        if dungeon.is_Calyx_Crimson:
+            return may_obtain_multi()
+
+        raise ScriptError(f'_obtain_get_entry: Cannot get entry from {dungeon}')
+
+    def _obtain_parse(self) -> ItemAmount | None:
+        """
+        Pages:
+            in: ITEM_CLOSE
+        """
+        ocr = OcrItemName(ITEM_NAME)
+        item = ocr.matched_single_line(self.device.image, keyword_classes=ITEM_CLASSES)
+        ocr = OcrItemAmount(ITEM_AMOUNT)
+        amount = ocr.ocr_single_line(self.device.image)
+
+        if item is None:
+            logger.warning('_obtain_parse: Unknown item name')
+            return None
+
+        logger.info(f'Item amount: item={item}, amount={amount}')
+        return ItemAmount(
+            item=item,
+            amount=amount,
+        )
+
+    def obtain_get(self, dungeon=None, skip_first_screenshot=True) -> list[ItemAmount]:
+        """
+        Args:
+            dungeon: Current dungeon,
+                or None for no early stop optimization
+            skip_first_screenshot:
+
+        Returns:
+            list[ItemAmount]:
+
+        Pages:
+            in: COMBAT_PREPARE
+            out: COMBAT_PREPARE
+        """
+        logger.hr('Obtain get', level=2)
+        if not skip_first_screenshot:
+            self.device.screenshot()
+
+        index = 1
+        prev = None
+        items = []
+        for _ in range(5):
+            entry = self._obtain_get_entry(dungeon, index=index, prev=prev)
+            if entry is None:
+                logger.info('Obtain get end')
+                break
+
+            self._obtain_enter(entry)
+            item = self._obtain_parse()
+            if item is not None:
+                items.append(item)
+                index += 1
+                prev = item
+            self._obtain_close()
+
+        logger.hr('Obtain get result')
+        for item in items:
+            logger.info(f'ItemAmount: {item.item.name}, {item.amount}')
+        """
+        <<< OBTAIN GET RESULT >>>                   
+        ItemAmount: Arrow_of_the_Starchaser, 15     
+        ItemAmount: Arrow_of_the_Demon_Slayer, 68   
+        ItemAmount: Arrow_of_the_Beast_Hunter, 85   
+        """
+        return items
+
+
+if __name__ == '__main__':
+    self = DungeonObtain('src')
+    self.device.screenshot()
+    self.obtain_get()

tasks/planner/keywords/classes.py

@@ -28,6 +28,22 @@ class ItemBase(Keyword):
         else:
             return None
 
+    @cached_property
+    def is_rarity_gold(self):
+        return self.rarity == 'SuperRare'
+
+    @cached_property
+    def is_rarity_purple(self):
+        return self.rarity == 'VeryRare'
+
+    @cached_property
+    def is_rarity_blue(self):
+        return self.rarity == 'Rare'
+
+    @cached_property
+    def is_rarity_green(self):
+        return self.rarity == 'NotNormal'
+
 
 @dataclass(repr=False)
 class ItemAscension(ItemBase):

tasks/planner/result.py

@@ -1,10 +1,12 @@
+import re
+
 import cv2
 from pponnxcr.predict_system import BoxedResult
 from pydantic import BaseModel
 
 from module.base.utils import area_center, area_in_area
 from module.logger import logger
-from module.ocr.ocr import OcrWhiteLetterOnComplexBackground
+from module.ocr.ocr import Ocr, OcrWhiteLetterOnComplexBackground
 from module.ui.scroll import AdaptiveScroll
 from tasks.daily.synthesize import SynthesizeUI
 from tasks.planner.assets.assets_planner_result import *
@@ -26,7 +28,16 @@ class PlannerResultRow(BaseModel):
         return self.item == other.item
 
 
-class OcrPlannerResult(OcrWhiteLetterOnComplexBackground):
+class OcrItemName(Ocr):
+    def after_process(self, result):
+        result = result.replace('念火之心', '忿火之心')
+        result = re.sub('工造机$', '工造机杼', result)
+        result = re.sub('工造轮', '工造迴轮', result)
+        result = re.sub('月狂牙', '月狂獠牙', result)
+        return result
+
+
+class OcrPlannerResult(OcrWhiteLetterOnComplexBackground, OcrItemName):
     def __init__(self):
         # Planner currently CN only
         super().__init__(OCR_RESULT, lang='cn')
@@ -146,7 +157,7 @@ class PlannerResult(SynthesizeUI):
         Pages:
             in: planner result
         """
-        logger.hr('Parse planner result')
+        logger.hr('Parse planner result', level=2)
         scroll = AdaptiveScroll(RESULT_SCROLL.button, name=RESULT_SCROLL.name)
         scroll.drag_threshold = 0.1
         scroll.edge_threshold = 0.1