From 8750704a4d1edf06a12665cc715599afea64202c Mon Sep 17 00:00:00 2001
From: LmeSzinc <37934724+LmeSzinc@users.noreply.github.com>
Date: Wed, 18 Oct 2023 21:51:51 +0800
Subject: [PATCH] Fix: Retry predict_door_by_name() to handle random OCR errors

---
 tasks/rogue/route/base.py |  2 +-
 tasks/rogue/route/exit.py | 19 +++++++++++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tasks/rogue/route/base.py b/tasks/rogue/route/base.py
index c4855072c..e144403e5 100644
--- a/tasks/rogue/route/base.py
+++ b/tasks/rogue/route/base.py
@@ -273,7 +273,7 @@ class RouteBase(RouteBase_, RogueExit, RogueEvent):
         self.rotation_set(end_rotation, threshold=10)
 
         logger.hr('Find domain exit', level=2)
-        direction = self.predict_door_by_name(self.device.image)
+        direction = self.predict_door(self.device.image)
         direction_limit = 55
         if direction is not None:
             if abs(direction) > direction_limit:
diff --git a/tasks/rogue/route/exit.py b/tasks/rogue/route/exit.py
index 39930ee19..c546128c6 100644
--- a/tasks/rogue/route/exit.py
+++ b/tasks/rogue/route/exit.py
@@ -1,5 +1,4 @@
 import re
-from typing import Optional
 
 import cv2
 import numpy as np
@@ -170,7 +169,7 @@ class RogueExit(CombatInteract):
         logger.info(f'PlanarDoor: {planar_door}, direction: {direction}')
         return direction
 
-    def predict_door_by_name(self, image) -> Optional[float]:
+    def predict_door_by_name(self, image) -> float | None:
         # Paint current name black
         x1, y1, x2, y2 = OCR_MAP_NAME.area
         image[y1:y2, x1:x2] = (0, 0, 0)
@@ -230,3 +229,19 @@ class RogueExit(CombatInteract):
         logger.error('No domain was selected, return the first instead')
         logger.info(f'Goto next domain: {results[0]}')
         return directions[0]
+
+    def predict_door(self, skip_first_screenshot=True) -> float | None:
+        timeout = Timer(3, count=6).start()
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if timeout.reached():
+                logger.error('Predict door timeout')
+                return None
+
+            direction = self.predict_door_by_name(self.device.image)
+            if direction is not None:
+                return direction