feat: implement robust screenshot handling with retry logic for image decoding

ba_auto/config.py

@@ -7,6 +7,14 @@ ENV = {**os.environ, "DISPLAY": DISPLAY, "XAUTHORITY": XAUTHORITY}
 
 WINDOW_NAME = "BlueArchive"
 
+# scrot occasionally writes a truncated/corrupt PNG (observed live as a
+# libpng "IDAT: invalid block type" decode error) even though the scrot
+# process itself exits 0 -- cv2.imread() returns None rather than raising on
+# a decode failure, so every color_at()/colors_at() call retries the whole
+# capture+decode cycle this many times before giving up, instead of crashing
+# the entire run on a single bad frame. See driver.py's _read_screenshot.
+SCREENSHOT_DECODE_RETRIES = 3
+
 # Resolved relative to wherever this checkout physically lives (two levels
 # up from ba_auto/config.py) rather than a fixed path outside the repo --
 # the runtime now runs directly out of the git checkout (no more separate

ba_auto/detector.py

@@ -61,9 +61,8 @@ def _masked_template(template):
 
 
 def find_cafe_sparkle():
-    driver.screenshot(SPARKLE_SHOT_PATH)
     template_full = cv2.imread(config.CAFE_SPARKLE_TEMPLATE)
-    img = cv2.imread(SPARKLE_SHOT_PATH)
+    img = driver.read_screenshot(SPARKLE_SHOT_PATH)
     th0, tw0 = template_full.shape[:2]
 
     best = None
@@ -103,8 +102,7 @@ def find_template(template_path, region=None, threshold=0.85):
     Returns the match's top-left `(x, y)` in full-screenshot coordinates, or
     None if nothing scores >= threshold.
     """
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     if region is not None:
         x1, y1, x2, y2 = region
         img = img[y1:y2, x1:x2]
@@ -132,8 +130,7 @@ def template_visible(template_path, region=None, threshold=0.85):
 
 def _color_mask(region, rgb_min, rgb_max):
     x1, y1, x2, y2 = region
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     b, g, r = crop[:, :, 0].astype(np.int16), crop[:, :, 1].astype(np.int16), crop[:, :, 2].astype(np.int16)
     (r_lo, g_lo, b_lo), (r_hi, g_hi, b_hi) = rgb_min, rgb_max
@@ -178,8 +175,7 @@ def find_color_centroid(region, rgb_min, rgb_max, min_pixels=1):
 
 def _ocr_crop(region):
     x1, y1, x2, y2 = region
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
     # This UI's text is consistently dark-on-light -- a hard black/white
@@ -222,8 +218,7 @@ def read_int_white_on_dark(region, psm=7):
     read_int_on_heart_badge's badge-outline contamination problem.
     """
     x1, y1, x2, y2 = region
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
     _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
@@ -262,8 +257,7 @@ def read_int_bordered(region, psm=7, border=20):
     margin.
     """
     x1, y1, x2, y2 = region
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
     _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
@@ -291,8 +285,7 @@ def read_int_on_heart_badge(region, psm=7):
     shape and keeps just the glyph.
     """
     x1, y1, x2, y2 = region
-    driver.screenshot(OCR_SHOT_PATH)
+    img = driver.read_screenshot(OCR_SHOT_PATH)
-    img = cv2.imread(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     b, g, r = crop[:, :, 0].astype(np.int16), crop[:, :, 1].astype(np.int16), crop[:, :, 2].astype(np.int16)
     ink = (r < g) & (np.maximum(np.maximum(r, g), b) < 170)

ba_auto/driver.py

@@ -98,9 +98,37 @@ def screenshot(path):
     run_command(["scrot", "-a", "0,0,1920,1200", "-o", path])
 
 
+def read_screenshot(path):
+    """screenshot() + cv2.imread(), retrying the whole capture if the decode
+    fails. scrot occasionally writes a truncated/corrupt PNG (observed live
+    as a libpng "IDAT: invalid block type" decode error) while still exiting
+    0, so a bad frame is a transient condition, not a hard failure -- same
+    "retry, don't crash" contract as the rest of this project's screen-state
+    checks. cv2.imread() returns None on a decode failure rather than
+    raising, which used to propagate straight into a bare image[y, x] and
+    crash the whole run (see plan.md/Handoff.md for the live incident this
+    was found from, mid schedule-grid-close in lesson.py).
+
+    Used for every screenshot-then-decode call site in this project
+    (color_at/colors_at below, and detector.py's OCR/template-match reads)
+    rather than each one pairing its own screenshot()+cv2.imread() and
+    duplicating this retry -- see detector.py's OCR_SHOT_PATH/
+    SPARKLE_SHOT_PATH call sites.
+    """
+    last_error = None
+    for attempt in range(1, config.SCREENSHOT_DECODE_RETRIES + 1):
+        screenshot(path)
+        image = cv2.imread(path)
+        if image is not None:
+            return image
+        last_error = f"cv2.imread returned None for {path} (attempt {attempt}/{config.SCREENSHOT_DECODE_RETRIES})"
+        print(f"[driver] {last_error} -- retrying capture")
+        wait(0.5)
+    raise RuntimeError(f"[driver] giving up on screenshot capture: {last_error}")
+
+
 def color_at(x, y):
-    screenshot(PROBE_SHOT_PATH)
+    image = read_screenshot(PROBE_SHOT_PATH)
-    image = cv2.imread(PROBE_SHOT_PATH)
     b, g, r = image[y, x]
     return int(r), int(g), int(b)
 
@@ -116,8 +144,7 @@ def colors_at(points):
 
     Returns a list of (r, g, b) tuples in the same order as `points`.
     """
-    screenshot(PROBE_SHOT_PATH)
+    image = read_screenshot(PROBE_SHOT_PATH)
-    image = cv2.imread(PROBE_SHOT_PATH)
     colors = []
     for x, y in points:
         b, g, r = image[y, x]