test-driver: User proper Path instead of str in OCR code
This commit is contained in:
parent
f56933ebbf
commit
e6ea13f4ea
@ -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
|
||||||
|
|
||||||
|
|||||||
@ -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],
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user