Add: Detect and switch characters

module/base/utils/utils.py

@@ -274,6 +274,20 @@ def area_size(area):
     )
 
 
+def area_center(area):
+    """
+    Get the center of an area
+
+    Args:
+        area: (upper_left_x, upper_left_y, bottom_right_x, bottom_right_y)
+
+    Returns:
+        tuple: (x, y)
+    """
+    x1, y1, x2, y2 = area
+    return (x1 + x2) / 2, (y1 + y2) / 2
+
+
 def point_limit(point, area):
     """
     Limit point in an area.

module/ocr/ocr.py

@@ -3,13 +3,14 @@ import time
 from datetime import timedelta
 from typing import Optional
 
+import cv2
 import numpy as np
 from pponnxcr.predict_system import BoxedResult
 
 import module.config.server as server
 from module.base.button import ButtonWrapper
 from module.base.decorator import cached_property
-from module.base.utils import area_pad, corner2area, crop, float2str
+from module.base.utils import area_pad, corner2area, crop, extract_white_letters, float2str
 from module.exception import ScriptError
 from module.logger import logger
 from module.ocr.keyword import Keyword
@@ -416,3 +417,20 @@ class Duration(Ocr):
         if number is None:
             return 0
         return int(number)
+
+
+class OcrWhiteLetterOnComplexBackground(Ocr):
+    def pre_process(self, image):
+        image = extract_white_letters(image, threshold=255)
+        image = cv2.merge([image, image, image])
+        return image
+
+    def detect_and_ocr(self, *args, **kwargs):
+        # Try hard to lower TextSystem.box_thresh
+        backup = self.model.text_detector.box_thresh
+        self.model.text_detector.box_thresh = 0.2
+
+        result = super().detect_and_ocr(*args, **kwargs)
+
+        self.model.text_detector.box_thresh = backup
+        return result

tasks/base/assets/assets_base_main_page.py

@@ -3,6 +3,16 @@ 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 ```
 
+OCR_CHARACTERS = ButtonWrapper(
+    name='OCR_CHARACTERS',
+    share=Button(
+        file='./assets/share/base/main_page/OCR_CHARACTERS.png',
+        area=(1041, 142, 1205, 436),
+        search=(1021, 122, 1225, 456),
+        color=(123, 130, 143),
+        button=(1041, 142, 1205, 436),
+    ),
+)
 OCR_MAP_NAME = ButtonWrapper(
     name='OCR_MAP_NAME',
     share=Button(

tasks/base/main_page.py

@@ -1,20 +1,26 @@
 import re
-from typing import Optional
+
+import cv2
+import numpy as np
+from scipy import signal
 
 import module.config.server as server
+from module.base.timer import Timer
+from module.base.utils import area_center, crop, rgb2luma
 from module.config.server import VALID_LANG
 from module.exception import RequestHumanTakeover, ScriptError
 from module.logger import logger
-from module.ocr.ocr import Ocr
-from tasks.base.assets.assets_base_main_page import OCR_MAP_NAME, ROGUE_LEAVE_FOR_NOW
+from module.ocr.ocr import OcrResultButton, OcrWhiteLetterOnComplexBackground
+from tasks.base.assets.assets_base_main_page import OCR_CHARACTERS, OCR_MAP_NAME, ROGUE_LEAVE_FOR_NOW
 from tasks.base.assets.assets_base_page import CLOSE, MAP_EXIT
 from tasks.base.page import Page, page_gacha, page_main
 from tasks.base.popup import PopupHandler
+from tasks.character.keywords import CharacterList
 from tasks.daily.assets.assets_daily_trial import START_TRIAL
 from tasks.map.keywords import KEYWORDS_MAP_PLANE, MapPlane
 
 
-class OcrPlaneName(Ocr):
+class OcrPlaneName(OcrWhiteLetterOnComplexBackground):
     def after_process(self, result):
         # RobotSettlement1
         result = re.sub(r'-[Ii1]$', '', result)
@@ -57,14 +63,31 @@ class OcrPlaneName(Ocr):
         return super().after_process(result)
 
 
+class OcrCharacterName(OcrWhiteLetterOnComplexBackground):
+    merge_thres_x = 20
+    merge_thres_y = 20
+
+    def after_process(self, result):
+        result = result.replace('蛆', '妲')
+
+        return super().after_process(result)
+
+
 class MainPage(PopupHandler):
     # Same as BigmapPlane class
     # Current plane
     plane: MapPlane = KEYWORDS_MAP_PLANE.Herta_ParlorCar
+    character_buttons: list[OcrResultButton] = []
+    character_current: CharacterList | None = None
 
     _lang_checked = False
 
-    def get_plane(self, lang=None) -> Optional[MapPlane]:
+    @property
+    def characters(self) -> list[CharacterList]:
+        characters = [button.matched_keyword for button in self.character_buttons]
+        return characters
+
+    def update_plane(self, lang=None) -> MapPlane | None:
         """
         Pages:
             in: page_main
@@ -89,7 +112,7 @@ class MainPage(PopupHandler):
 
         return None
 
-    def check_lang_from_map_plane(self) -> Optional[str]:
+    def check_lang_from_map_plane(self) -> str | None:
         logger.info('check_lang_from_map_plane')
         lang_unknown = self.config.Emulator_GameLanguage == 'auto'
 
@@ -101,7 +124,7 @@ class MainPage(PopupHandler):
 
         for lang in lang_list:
             logger.info(f'Try ocr in lang {lang}')
-            keyword = self.get_plane(lang)
+            keyword = self.update_plane(lang)
             if keyword is not None:
                 logger.info(f'check_lang_from_map_plane matched lang: {lang}')
                 if lang_unknown or lang != server.lang:
@@ -151,6 +174,140 @@ class MainPage(PopupHandler):
         self.handle_lang_check(page=page_main)
         return True
 
+    def update_characters(self) -> list[CharacterList]:
+        ocr = OcrCharacterName(OCR_CHARACTERS)
+        self.character_buttons = ocr.matched_ocr(self.device.image, keyword_classes=CharacterList)
+        characters = self.characters
+        logger.attr('Characters', characters)
+        self.character_current = self._convert_selected_to_character(self._update_current_character())
+        return characters
+
+    def _update_current_character(self) -> list[int]:
+        """
+        Returns:
+            list[int]: Selected index, 1 to 4.
+        """
+        # 50px-width area starting from the right edge of HP bars
+        area = (1101, 151, 1151, 459)
+        # Y coordinates where the color peaks should be when character is selected
+        expected_peaks = np.array([201, 279, 357, 435])
+        expected_peaks_in_area = expected_peaks - area[1]
+        # Use Luminance to fit H264 video stream
+        image = rgb2luma(crop(self.device.image, area))
+        # Remove character names
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+        image = cv2.erode(image, kernel)
+        # To find peaks along Y
+        line = cv2.reduce(image, 1, cv2.REDUCE_AVG).flatten().astype(int)
+
+        # Find color peaks
+        parameters = {
+            'height': (60, 255),
+            'prominence': 30,
+            'distance': 5,
+        }
+        peaks, _ = signal.find_peaks(line, **parameters)
+        # Remove smooth peaks
+        parameters = {
+            'height': (5, 255),
+            'prominence': 5,
+            'distance': 5,
+        }
+        diff = -np.diff(line)
+        diff_peaks, _ = signal.find_peaks(diff, **parameters)
+
+        def is_steep_peak(y, threshold=5):
+            return np.abs(diff_peaks - y).min() <= threshold
+
+        def peak_to_selected(y, threshold=5):
+            distance = np.abs(expected_peaks_in_area - y)
+            return np.argmin(distance) + 1 if distance.min() < threshold else 0
+
+        selected = [peak_to_selected(peak) for peak in peaks if peak_to_selected(peak) and is_steep_peak(peak)]
+        logger.attr('CharacterSelected', selected)
+        return selected
+
+    def _convert_selected_to_character(self, selected: list[int]) -> CharacterList | None:
+        expected_peaks = np.array([201, 279, 357, 435])
+        if not selected:
+            logger.warning(f'No current character')
+            logger.attr('CurrentCharacter', None)
+            return None
+        elif len(selected) == 1:
+            selected = selected[0]
+        else:
+            logger.warning(f'Too many current characters: {selected}, using first')
+            selected = selected[0]
+
+        expected_y = expected_peaks[selected - 1]
+        for button in self.character_buttons:
+            y = area_center(button.area)[1]
+            if expected_y - 78 < y < expected_y:
+                logger.attr('CurrentCharacter', button.matched_keyword)
+                return button.matched_keyword
+
+        logger.warning(f'Current character: {selected} does not belong to any detected character')
+        logger.attr('CurrentCharacter', None)
+        return None
+
+    def character_switch(self, character: CharacterList | str | int, skip_first_screenshot=True) -> bool:
+        """
+        Args:
+            character: CharacterList object, or character name, or select index from 1 to 4.
+            skip_first_screenshot:
+
+        Returns:
+            bool: If chose
+        """
+        logger.info(f'Character choose: {character}')
+        characters = self.characters
+        if isinstance(character, int):
+            character = self._convert_selected_to_character([character])
+            if character is None:
+                return False
+            try:
+                index = characters.index(character) + 1
+            except IndexError:
+                logger.warning(f'Cannot choose character {character} as it was not detected')
+                return False
+        else:
+            if isinstance(character, str):
+                character = CharacterList.find(character)
+            try:
+                index = characters.index(character) + 1
+            except IndexError:
+                logger.warning(f'Cannot choose character {character} as it was not detected')
+                return False
+
+        button = self.character_buttons[index - 1]
+        interval = Timer(1, count=3)
+        count = 0
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            # End
+            selected = self._update_current_character()
+            if index in selected:
+                logger.info('Character chose')
+                return True
+            if count > 3:
+                logger.warning('Failed to choose character, assume chose')
+                return False
+
+            try:
+                is_in_main = self.is_in_main
+            except AttributeError:
+                logger.critical('Method ui_goto() not found, class MainPage must be inherited by class UI')
+                raise ScriptError
+
+            if interval.reached() and is_in_main():
+                self.device.click(button)
+                interval.reset()
+                count += 1
+
     def ui_leave_special(self):
         """
         Leave from:

tasks/rogue/route/exit.py

@@ -1,10 +1,9 @@
 import re
 
-import cv2
 import numpy as np
 
 from module.base.timer import Timer
-from module.base.utils import Points, extract_white_letters
+from module.base.utils import Points, area_center
 from module.logger import logger
 from tasks.base.assets.assets_base_main_page import OCR_MAP_NAME
 from tasks.base.main_page import OcrPlaneName
@@ -16,38 +15,9 @@ from tasks.rogue.assets.assets_rogue_ui import BLESSING_CONFIRM
 from tasks.rogue.assets.assets_rogue_weekly import ROGUE_REPORT
 
 
-def area_center(area):
-    """
-    Get the center of an area
-
-    Args:
-        area: (upper_left_x, upper_left_y, bottom_right_x, bottom_right_y)
-
-    Returns:
-        tuple: (x, y)
-    """
-    x1, y1, x2, y2 = area
-    return (x1 + x2) / 2, (y1 + y2) / 2
-
-
 class OcrDomainExit(OcrPlaneName):
     merge_thres_x = 50
 
-    def pre_process(self, image):
-        image = extract_white_letters(image, threshold=255)
-        image = cv2.merge([image, image, image])
-        return image
-
-    def detect_and_ocr(self, *args, **kwargs):
-        # Try hard to lower TextSystem.box_thresh
-        backup = self.model.text_detector.box_thresh
-        self.model.text_detector.box_thresh = 0.2
-
-        result = super().detect_and_ocr(*args, **kwargs)
-
-        self.model.text_detector.box_thresh = backup
-        return result
-
     def _match_result(
             self,
             result: str,

tasks/rogue/route/loader.py

@@ -81,7 +81,7 @@ class RouteLoader(RogueUI, MinimapWrapper, RouteLoader_, MainPage):
         Try to find from known route spawn point
         """
         logger.info('position_find_known')
-        plane = self.get_plane()
+        plane = self.update_plane()
         if plane is None:
             logger.warning('Unknown rogue domain')
             return