mirror of
https://github.com/LmeSzinc/StarRailCopilot.git
synced 2024-11-23 00:52:22 +00:00
164 lines
5.5 KiB
Python
164 lines
5.5 KiB
Python
import re
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from module.base.timer import Timer
|
|
from module.base.utils import crop
|
|
from module.logger import logger
|
|
from module.ocr.ocr import Digit, DigitCounter
|
|
from tasks.base.ui import UI
|
|
from tasks.combat.assets.assets_combat_stamina_status import *
|
|
|
|
|
|
class StaminaOcr(DigitCounter):
|
|
def after_process(self, result):
|
|
result = super().after_process(result)
|
|
# The trailblaze power icon is recognized as 买
|
|
# OCR_TRAILBLAZE_POWER includes the icon because the length varies by value
|
|
result = re.sub(r'[买米装来:()]', '', result)
|
|
# 61240 -> 6/240
|
|
result = re.sub(r'1240$', '/240', result)
|
|
# 0*0/24 -> 0/240
|
|
result = re.sub(r'24$', '240', result)
|
|
# * 50/2401+)
|
|
result = result.replace('/2401', '/240')
|
|
return result
|
|
|
|
|
|
class ReservedOcr(Digit):
|
|
pass
|
|
|
|
|
|
class ImmersifierOcr(DigitCounter):
|
|
pass
|
|
|
|
|
|
class DataStaminaStatus(BaseModel):
|
|
stamina: int | None
|
|
reserved: int | None
|
|
immersifier: int | None
|
|
|
|
|
|
class StaminaStatus(UI):
|
|
@staticmethod
|
|
def get_stamina_status(image) -> DataStaminaStatus:
|
|
"""
|
|
Update trailblaze power, stored trailblaze power, immersifier
|
|
|
|
Returns:
|
|
int: Stamina, or None if stamina not displayed or error on OCR
|
|
int: Reserved stamina
|
|
int: Immersifier
|
|
"""
|
|
for button in [STAMINA_ICON, RESERVED_ICON, IMMERSIFIER_ICON]:
|
|
button.load_search(ICON_SEARCH.area)
|
|
|
|
stamina = None
|
|
if STAMINA_ICON.match_template(image):
|
|
STAMINA_OCR.load_offset(STAMINA_ICON)
|
|
im = crop(image, STAMINA_OCR.button, copy=False)
|
|
stamina, _, total = StaminaOcr(STAMINA_OCR).ocr_single_line(im, direct_ocr=True)
|
|
if total > 240 or total == 0:
|
|
logger.warning(f'Unexpected stamina total: {total}')
|
|
stamina = None
|
|
|
|
reserved = None
|
|
if RESERVED_ICON.match_template(image):
|
|
RESERVED_OCR.load_offset(RESERVED_ICON)
|
|
im = crop(image, RESERVED_OCR.button, copy=False)
|
|
reserved = ReservedOcr(RESERVED_OCR).ocr_single_line(im, direct_ocr=True)
|
|
if reserved > 2400:
|
|
logger.warning(f'Unexpected reserved value: {reserved}')
|
|
reserved = None
|
|
|
|
immersifier = None
|
|
if IMMERSIFIER_ICON.match_template(image):
|
|
IMMERSIFIER_OCR.load_offset(IMMERSIFIER_ICON)
|
|
im = crop(image, IMMERSIFIER_OCR.button, copy=False)
|
|
immersifier, _, total = StaminaOcr(IMMERSIFIER_OCR).ocr_single_line(im, direct_ocr=True)
|
|
if total != 8:
|
|
logger.warning(f'Unexpected immersifier total: {total}')
|
|
immersifier = None
|
|
|
|
return DataStaminaStatus(
|
|
stamina=stamina,
|
|
reserved=reserved,
|
|
immersifier=immersifier,
|
|
)
|
|
|
|
def update_stamina_status(
|
|
self,
|
|
image=None,
|
|
skip_first_screenshot=True,
|
|
expect_stamina=False,
|
|
expect_reserved=False,
|
|
expect_immersifier=False,
|
|
) -> DataStaminaStatus:
|
|
"""
|
|
Update stamina status with retry
|
|
|
|
Args:
|
|
image: Detect given image only, no new screenshot will be taken
|
|
and all expect_* are considered False
|
|
skip_first_screenshot:
|
|
expect_stamina:
|
|
True to expect stamina exists, retry detect if it wasn't
|
|
expect_reserved:
|
|
expect_immersifier:
|
|
|
|
Pages:
|
|
in: page_guild, Survival_Index, Simulated_Universe
|
|
or page_rogue, LEVEL_CONFIRM
|
|
or rogue, REWARD_CLOSE
|
|
"""
|
|
timeout = Timer(1, count=2).start()
|
|
if image is None:
|
|
image = self.device.image
|
|
use_cached_image = False
|
|
else:
|
|
skip_first_screenshot = True
|
|
use_cached_image = True
|
|
|
|
while 1:
|
|
if skip_first_screenshot:
|
|
skip_first_screenshot = False
|
|
else:
|
|
self.device.screenshot()
|
|
image = self.device.image
|
|
|
|
# Timeout
|
|
if timeout.reached():
|
|
logger.warning('dungeon_update_stamina() timeout')
|
|
return DataStaminaStatus(
|
|
stamina=None,
|
|
reserved=None,
|
|
immersifier=None,
|
|
)
|
|
|
|
# Ocr
|
|
status = self.get_stamina_status(image)
|
|
valid = True
|
|
if expect_stamina and status.stamina is None:
|
|
valid = False
|
|
if expect_reserved and status.reserved is None:
|
|
valid = False
|
|
if expect_immersifier and status.immersifier is None:
|
|
valid = False
|
|
if status.stamina is None and status.reserved is None and status.immersifier is None:
|
|
logger.warning('update_stamina_status: No icon detected')
|
|
valid = False
|
|
|
|
# Write config
|
|
with self.config.multi_set():
|
|
if status.stamina is not None:
|
|
self.config.stored.TrailblazePower.value = status.stamina
|
|
if status.reserved is not None:
|
|
self.config.stored.Reserved.value = status.reserved
|
|
if status.immersifier is not None:
|
|
self.config.stored.Immersifier.value = status.immersifier
|
|
|
|
if use_cached_image or valid:
|
|
return status
|
|
else:
|
|
continue
|