test-driver: User proper Path instead of str in OCR code

This commit is contained in:
Jacek Galowicz 2025-07-01 13:05:00 +02:00
parent f56933ebbf
commit e6ea13f4ea
2 changed files with 10 additions and 9 deletions

View File

@ -834,13 +834,13 @@ class Machine:
self.connected = True self.connected = True
@contextmanager @contextmanager
def _managed_screenshot(self) -> Generator[str]: def _managed_screenshot(self) -> Generator[Path]:
""" """
Take a screenshot and yield the screenshot filepath. Take a screenshot and yield the screenshot filepath.
The file will be deleted when leaving the generator. The file will be deleted when leaving the generator.
""" """
with tempfile.TemporaryDirectory() as tmpdir: with tempfile.TemporaryDirectory() as tmpdir:
screenshot_path: str = os.path.join(tmpdir, "ppm") screenshot_path: Path = Path(tmpdir) / "ppm"
self.send_monitor_command(f"screendump {screenshot_path}") self.send_monitor_command(f"screendump {screenshot_path}")
yield screenshot_path yield screenshot_path

View File

@ -2,11 +2,12 @@ import multiprocessing
import os import os
import shutil import shutil
import subprocess import subprocess
from pathlib import Path
from test_driver.errors import MachineError from test_driver.errors import MachineError
def perform_ocr_on_screenshot(screenshot_path: str) -> str: def perform_ocr_on_screenshot(screenshot_path: Path) -> str:
""" """
Perform OCR on a screenshot that contains text. Perform OCR on a screenshot that contains text.
Returns a string with all words that could be found. Returns a string with all words that could be found.
@ -18,7 +19,7 @@ def perform_ocr_on_screenshot(screenshot_path: str) -> str:
def perform_ocr_variants_on_screenshot( def perform_ocr_variants_on_screenshot(
screenshot_path: str, variants: bool = True screenshot_path: Path, variants: bool = True
) -> list[str]: ) -> list[str]:
""" """
Same as perform_ocr_on_screenshot but will create variants of the images Same as perform_ocr_on_screenshot but will create variants of the images
@ -34,7 +35,7 @@ def perform_ocr_variants_on_screenshot(
# https://github.com/tesseract-ocr/tesseract/issues/3109 # https://github.com/tesseract-ocr/tesseract/issues/3109
processes = max(1, int(os.process_cpu_count() / 4)) processes = max(1, int(os.process_cpu_count() / 4))
with multiprocessing.Pool(processes=processes) as pool: with multiprocessing.Pool(processes=processes) as pool:
image_paths: list[str] = [screenshot_path] image_paths: list[Path] = [screenshot_path]
if variants: if variants:
image_paths.extend( image_paths.extend(
pool.starmap( pool.starmap(
@ -45,7 +46,7 @@ def perform_ocr_variants_on_screenshot(
return pool.map(_run_tesseract, image_paths) return pool.map(_run_tesseract, image_paths)
def _run_tesseract(image: str) -> str: def _run_tesseract(image: Path) -> str:
# tesseract --help-oem # tesseract --help-oem
# OCR Engine modes (OEM): # OCR Engine modes (OEM):
# 0|tesseract_only Legacy engine only. # 0|tesseract_only Legacy engine only.
@ -73,7 +74,7 @@ def _run_tesseract(image: str) -> str:
return ret.stdout.decode("utf-8") return ret.stdout.decode("utf-8")
def _preprocess_screenshot(screenshot_path: str, negate: bool = False) -> str: def _preprocess_screenshot(screenshot_path: Path, negate: bool = False) -> Path:
if shutil.which("magick") is None: if shutil.which("magick") is None:
raise MachineError("OCR requested but `magick` is not available") raise MachineError("OCR requested but `magick` is not available")
@ -98,7 +99,7 @@ def _preprocess_screenshot(screenshot_path: str, negate: bool = False) -> str:
if negate: if negate:
magick_args.append("-negate") magick_args.append("-negate")
out_file += ".negative" out_file = out_file.with_suffix(".negative")
magick_args += [ magick_args += [
"-gamma", "-gamma",
@ -106,7 +107,7 @@ def _preprocess_screenshot(screenshot_path: str, negate: bool = False) -> str:
"-blur", "-blur",
"1x65535", "1x65535",
] ]
out_file += ".png" out_file = out_file.with_suffix(".png")
ret = subprocess.run( ret = subprocess.run(
["magick", "convert"] + magick_args + [screenshot_path, out_file], ["magick", "convert"] + magick_args + [screenshot_path, out_file],